{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "578f9d15-3df9-461f-befc-ab77c381a96d",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "(Intro_PyTensor)=\n",
    "\n",
    "# What is PyTensor?\n",
    ":::{post} August 16, 2025 \n",
    ":tags: introduction, worked examples, tutorial\n",
    ":category: beginner, explanation \n",
    ":author: Jesse Grabowski, Ricardo Vieira\n",
    ":::\n",
    "\n",
    "A library to define, manipulate, and compile computational graphs.\n",
    "\n",
    "\n",
    "## Let's break it apart\n",
    "A library to (1.) define, (2.) manipulate, and (3.) compile (0.) computational graphs."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96f0aba6",
   "metadata": {},
   "source": [
    "## Accompanying Video\n",
    "\n",
    "The following video is a recording of a talk presenting the content of this notebook.\n",
    "\n",
    "<center>\n",
    "<iframe width=\"800\" height=\"450\" src=\"https://www.youtube.com/embed/C1JJ33z9g00\" title=\"YouTube video player\" frameborder=\"0\" allow=\"accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture; web-share\" allowfullscreen></iframe>\n",
    "</center>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0427f803",
   "metadata": {},
   "source": [
    "## Prepare Notebook\n",
    "\n",
    ":::{include} ../extra_installs.md\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "701a0893",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "plt.style.use(\"seaborn-v0_8\")\n",
    "\n",
    "%config InlineBackend.figure_format = \"retina\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdf5f119-c085-48a1-b6b0-733ed4921303",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## (0.) Computational graph\n",
    "\n",
    "Any program implies a computational graph. In PyTensor we're mostly focusing on static, array-based (i.e, numpy) programs with some branching and looping primitives. PyTensor is hackable, and can be easily extended to represent arbitrary types and operations. That said, its usefulness quickly vanishes as you venture out of its area of focus.\n",
    "\n",
    "Everyone here is likely familar with the idea of a computation graph, but let's look at a quick example anyway. Consider a program that computes $z = x + y$. Here we have:\n",
    "\n",
    "- Two symbolic inputs, $x$ and $y$\n",
    "- An operator, $+$, that takes several inputs and maps them to a single output\n",
    "- A symbolic output, $z$\n",
    "\n",
    "Visualized:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f1890c77",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "import graphviz as gr\n",
    "\n",
    "\n",
    "def draw_graph(edge_list, node_props=None, edge_props=None, graph_direction=\"UD\"):\n",
    "    \"\"\"Utility to draw a causal (directed) graph\"\"\"\n",
    "    g = gr.Digraph(\n",
    "        graph_attr={\n",
    "            \"rankdir\": graph_direction,\n",
    "            \"ratio\": \"0.3\",\n",
    "            \"overlap\": \"vpsc\",\n",
    "            \"splines\": \"true\",\n",
    "            \"mode\": \"sgd\",\n",
    "            \"lheight\": \"4\",\n",
    "        },\n",
    "        engine=\"dot\",\n",
    "    )\n",
    "\n",
    "    edge_props = {} if edge_props is None else edge_props\n",
    "    for e in edge_list:\n",
    "        props = edge_props[e] if e in edge_props else {}\n",
    "        g.edge(e[0], e[1], **props)\n",
    "\n",
    "    if node_props is not None:\n",
    "        for name, props in node_props.items():\n",
    "            g.node(name=name, **props)\n",
    "    return g\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "870d7457",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "While this DAG is boring and obvious, it highlights that there are two types of nodes:\n",
    "\n",
    "- Variables\n",
    "- Operations\n",
    "\n",
    "We can also distinguish between \"root variables\", like $x$ and $y$, intermediate variables (we don't have any here :( ) and output variables ($z$). These differences will return in a few slides!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e0396698",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 13.0.0 (20250608.1624)\n",
       " -->\n",
       "<!-- Pages: 1 -->\n",
       "<svg width=\"608pt\" height=\"188pt\"\n",
       " viewBox=\"0.00 0.00 608.00 188.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 184)\">\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-184 604,-184 604,4 -4,4\"/>\n",
       "<!-- x -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>x</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"129\" cy=\"-162\" rx=\"27\" ry=\"18\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"129\" y=\"-156.95\" font-family=\"Times,serif\" font-size=\"14.00\">x</text>\n",
       "</g>\n",
       "<!-- + -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>+</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"300\" cy=\"-90\" rx=\"27\" ry=\"18\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"300\" y=\"-84.95\" font-family=\"Times,serif\" font-size=\"14.00\">+</text>\n",
       "</g>\n",
       "<!-- x&#45;&gt;+ -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>x&#45;&gt;+</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M151.71,-151.7C181.22,-139.62 232.98,-118.44 266.87,-104.56\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"267.87,-107.93 275.8,-100.91 265.22,-101.46 267.87,-107.93\"/>\n",
       "</g>\n",
       "<!-- z -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>z</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"300\" cy=\"-18\" rx=\"27\" ry=\"18\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"300\" y=\"-12.95\" font-family=\"Times,serif\" font-size=\"14.00\">z</text>\n",
       "</g>\n",
       "<!-- +&#45;&gt;z -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>+&#45;&gt;z</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M300,-71.7C300,-64.41 300,-55.73 300,-47.54\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"303.5,-47.62 300,-37.62 296.5,-47.62 303.5,-47.62\"/>\n",
       "</g>\n",
       "<!-- y -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>y</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"471\" cy=\"-162\" rx=\"27\" ry=\"18\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"471\" y=\"-156.95\" font-family=\"Times,serif\" font-size=\"14.00\">y</text>\n",
       "</g>\n",
       "<!-- y&#45;&gt;+ -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>y&#45;&gt;+</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M448.29,-151.7C418.78,-139.62 367.02,-118.44 333.13,-104.56\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"334.78,-101.46 324.2,-100.91 332.13,-107.93 334.78,-101.46\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x11f29e690>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "draw_graph([(\"x\", \"+\"), (\"y\", \"+\"), (\"+\", \"z\")])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97944cd3",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## (1.) Definition \n",
    "In PyTensor, you define a computational graph explicitly, starting with _placeholder_ input variables. From these inputs you build more intermediate variables by applying operators (like $+$), which can then be treated either as outputs, or as intermediate variables for further computation.\n",
    "\n",
    "To reduce the learning barrier, Pytensor was designed to look like numpy code. But be aware that it's not! Let' look at some differences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ed6d4a86d7398eca",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.        , 0.69314718, 1.31326169])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "\n",
    "# Numpy\n",
    "x = np.array([0, 1, np.e])  # Actual numbers\n",
    "y = np.log(1 + x)  # Actual computation\n",
    "y  # Actual result"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "559a56aa",
   "metadata": {},
   "source": [
    "Now, let's see how we can use PyTensor to define a computational graph for the same operation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f10206d3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Log.0, pytensor.tensor.variable.TensorVariable)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pytensor\n",
    "import pytensor.tensor as pt\n",
    "\n",
    "\n",
    "# Pytensor\n",
    "x = pt.tensor(shape=(3,), dtype=\"float64\")  # Symbolic vector\n",
    "y = pt.log(1 + x)  # Symbolic computation\n",
    "y, type(y)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd6ebfa4",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "The most important thing to grok here is that `y` is not a number! It's a symbol. Specifically, it's the output of a `Log` operation, so we see the name `Log.0`. The type is a `TensorVariable`, which is a basic unit of symbolic computation.\n",
    "\n",
    "Notice how Pytensor is straddling a line between a computer algebra system like sympy or maple, and a tensor library like numpy."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "904410ca",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "A key tool when working with Pytensor is  {func}`~pytensor.dprint`. It's short for \"debug print\", and it shows you a text representation of a graph. It's always a good idea to look at the graph pytensor is generating for your graph, because it shows you exactly what is going on. Admittedly, they take a bit of getting used to reading."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f6f5731649ea29c7",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Log [id A]\n",
      " └─ Add [id B]\n",
      "    ├─ ExpandDims{axis=0} [id C]\n",
      "    │  └─ 1 [id D]\n",
      "    └─ <Vector(float64, shape=(3,))> [id E]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x105c9e0e0>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y.dprint()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e913e973",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "PyTensor can also generate visual representations of your graph. It's great, but doesn't scale well.\n",
    "\n",
    "That said, here's the visual representation of our graph. It's very similar to our $z = x + y$ graph, with a some extra computational accoutrement. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "25e22315",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<svg xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\" width=\"378pt\" height=\"367pt\" viewBox=\"0.00 0.00 378.00 367.00\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 363)\">\n",
       "<title>G</title>\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-363 374.08,-363 374.08,4 -4,4\"/>\n",
       "<!-- 5031590032 -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>5031590032</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"276.25\" cy=\"-268\" rx=\"93.83\" ry=\"18\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"276.25\" y=\"-262.95\" font-family=\"Times,serif\" font-size=\"14.00\">ExpandDims{axis=0}</text>\n",
       "</g>\n",
       "<!-- 4817926624 -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>4817926624</title>\n",
       "<ellipse fill=\"#ffaabb\" stroke=\"black\" cx=\"97.25\" cy=\"-179.5\" rx=\"27\" ry=\"18\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"97.25\" y=\"-174.45\" font-family=\"Times,serif\" font-size=\"14.00\">Add</text>\n",
       "</g>\n",
       "<!-- 5031590032&#45;&gt;4817926624 -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>5031590032-&gt;4817926624</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M242.58,-250.73C209.92,-234.95 160.73,-211.18 128.8,-195.75\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"130.58,-192.72 120.05,-191.52 127.53,-199.02 130.58,-192.72\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"273.28\" y=\"-218.7\" font-family=\"Times,serif\" font-size=\"14.00\">0 Vector(int8, shape=(1,))</text>\n",
       "</g>\n",
       "<!-- 4819868256 -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>4819868256</title>\n",
       "<polygon fill=\"green\" stroke=\"black\" points=\"361.5,-359 191,-359 191,-323 361.5,-323 361.5,-359\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"276.25\" y=\"-335.95\" font-family=\"Times,serif\" font-size=\"14.00\">val=1 Scalar(int8, shape=())</text>\n",
       "</g>\n",
       "<!-- 4819868256&#45;&gt;5031590032 -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>4819868256-&gt;5031590032</title>\n",
       "<path fill=\"none\" stroke=\"blue\" d=\"M276.25,-322.81C276.25,-315.23 276.25,-306.1 276.25,-297.54\"/>\n",
       "<polygon fill=\"blue\" stroke=\"blue\" points=\"279.75,-297.54 276.25,-287.54 272.75,-297.54 279.75,-297.54\"/>\n",
       "</g>\n",
       "<!-- 5036460800 -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>5036460800</title>\n",
       "<ellipse fill=\"#ffaabb\" stroke=\"black\" cx=\"97.25\" cy=\"-91\" rx=\"27\" ry=\"18\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"97.25\" y=\"-85.95\" font-family=\"Times,serif\" font-size=\"14.00\">Log</text>\n",
       "</g>\n",
       "<!-- 4817926624&#45;&gt;5036460800 -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>4817926624-&gt;5036460800</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M97.25,-161.41C97.25,-149.76 97.25,-134.05 97.25,-120.52\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"100.75,-120.86 97.25,-110.86 93.75,-120.86 100.75,-120.86\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"171.5\" y=\"-130.2\" font-family=\"Times,serif\" font-size=\"14.00\">Vector(float64, shape=(3,))</text>\n",
       "</g>\n",
       "<!-- 4990476096 -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>4990476096</title>\n",
       "<polygon fill=\"green\" stroke=\"black\" points=\"164.5,-286 0,-286 0,-250 164.5,-250 164.5,-286\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"82.25\" y=\"-262.95\" font-family=\"Times,serif\" font-size=\"14.00\">Vector(float64, shape=(3,))</text>\n",
       "</g>\n",
       "<!-- 4990476096&#45;&gt;4817926624 -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>4990476096-&gt;4817926624</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M85.21,-249.91C87.23,-238.26 89.96,-222.55 92.3,-209.02\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"95.72,-209.79 93.98,-199.34 88.83,-208.6 95.72,-209.79\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"94.47\" y=\"-218.7\" font-family=\"Times,serif\" font-size=\"14.00\">1</text>\n",
       "</g>\n",
       "<!-- 5035578032 -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>5035578032</title>\n",
       "<polygon fill=\"blue\" stroke=\"black\" points=\"179.5,-36 15,-36 15,0 179.5,0 179.5,-36\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"97.25\" y=\"-12.95\" font-family=\"Times,serif\" font-size=\"14.00\">Vector(float64, shape=(3,))</text>\n",
       "</g>\n",
       "<!-- 5036460800&#45;&gt;5035578032 -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>5036460800-&gt;5035578032</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M97.25,-72.81C97.25,-65.23 97.25,-56.1 97.25,-47.54\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"100.75,-47.54 97.25,-37.54 93.75,-47.54 100.75,-47.54\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>"
      ],
      "text/plain": [
       "<IPython.core.display.SVG object>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython.display import SVG\n",
    "\n",
    "from pytensor.printing import pydotprint\n",
    "\n",
    "\n",
    "SVG(pydotprint(y, return_image=True, format=\"svg\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ec992ea",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "We can see these definitions in action by looking at properties of `y`. We've already seen that `type(y)` is a `TensorVariable`. With the exception of root variables, all `TensorVariables` have an `owner`, which is the `Apply` node that created it. \n",
    "\n",
    "In this case, the owner of `y` is `Log(Add.0)`. `Add.0` is the name of the output of an `Add` `Op`, so `y` is the result of a `Log` applied to an `Add`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2c0a5be5591f936",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Log(Add.0), pytensor.graph.basic.Apply)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y.owner, type(y.owner)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ef053ec",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "In this case, the specific Op is an `Elemwise`, which is a meta-Op that broadcasts a scalar computation (in this case `scalar_op = log`) across a tensor input. It makes it an elementwise operation... hence \"Elemwise Op\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "bc75e500fb9ece2f",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Elemwise(scalar_op=log,inplace_pattern=<frozendict {}>),\n",
       " pytensor.tensor.elemwise.Elemwise)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y.owner.op, type(y.owner.op)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e2e3fac",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "We can also check the outputs of the `Apply`. This is, of course, `y` itself!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f1bb3439391f2509",
   "metadata": {
    "scrolled": true,
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([Log.0], True)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y.owner.outputs, y.owner.outputs == [y]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75685087",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Finally, there's the inputs. This is `Add.0`, which represents `1 + x`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c6f97556431c818f",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Add.0]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y.owner.inputs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b57e98dc",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "From here, the story begin again, and we can keep climbing up the graph until the root variable, `x`. I'll spare you, though"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "316c181a13627039",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Add [id A]\n",
      " \u251c\u2500 ExpandDims{axis=0} [id B]\n",
      " \u2502  \u2514\u2500 1 [id C]\n",
      " \u2514\u2500 <Vector(float64, shape=(3,))> [id D]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x105c9e0e0>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y.owner.inputs[0].dprint()  # And the story begins again"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72133ce2",
   "metadata": {},
   "source": [
    "```{note}\n",
    "For those curious: This kind of graph is a bi-partite, directed, acyclic graph composed of interconnected Variable -> Apply -> Variable nodes.\n",
    "\n",
    "Apply nodes connect input variables to output variables, via a specific operator. Variables have a type and can have an owner (the Apply node that creates them) or not (if they are root placeholder variables).\n",
    "\n",
    "Here is a schematic of the $z = x + y$ graph again, but this time with annotations that show how it is represented in Pytensor. I always felt it was upside-down; perhaps that feeling will help you to understand what you're looking at. Read it from top-to-bottom, but with the flow of computation running against the arrows:\n",
    "\n",
    "![](https://pytensor.readthedocs.io/en/latest/_images/apply.png)\n",
    "\n",
    "Looking at the schematic again, note that there is a conceptual difference between `Apply` (which is a big box that handles a bunch of stuff) and an `Op`, which is the actual type of computation being done. \n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a916500-4dfe-4ef0-90ac-84f5de9f7ce2",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## (2) Manipulation\n",
    "\n",
    "So far, nothing we've seen is special. Maybe it's a nice API for writing functions, but it's not at all clear why the program `y = pt.log(1 + x)` is better than the following Python program:\n",
    "\n",
    "```py\n",
    "def log1p(x):\n",
    "    return np.log(1 + x)\n",
    "    \n",
    "y = log1p(x)\n",
    "```\n",
    "\n",
    "The answer is that PyTensor is able to **manipulate** the computation graph.\n",
    "\n",
    "There are three important graph operations to consider:\n",
    "\n",
    "1. Replacement\n",
    "2. Rewriting\n",
    "3. Transformation\n",
    "\n",
    "We will look at each of these in turn."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ec7dfd0",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Graph Replacement\n",
    "\n",
    "The simplest graph operation is a **replace**. We simply take one operation and swap it out for another. Consider the following graph:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "70351014",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Add [id A]\n",
      " ├─ ExpandDims{axis=0} [id B]\n",
      " │  └─ 1 [id C]\n",
      " └─ Sin [id D]\n",
      "    └─ Mul [id E]\n",
      "       ├─ ExpandDims{axis=0} [id F]\n",
      "       │  └─ 6.283185307179586 [id G]\n",
      "       └─ b [id H]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x105c9e0e0>"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a = pt.tensor(\"b\", shape=(None,))\n",
    "b = 2 * pt.pi * a\n",
    "c = pt.sin(b)\n",
    "d = 1 + c\n",
    "d.dprint()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87f8371c",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "If, for some reason, we wanted to do replace `pt.sin(b)` with `pt.cos(b)`, we could do so with the `graph_replace` function. The subgraph we want to target is `c`, and the top-level output is `d`, so we do it as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "a20df8d6",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Add [id A]\n",
      " ├─ ExpandDims{axis=0} [id B]\n",
      " │  └─ 1 [id C]\n",
      " └─ Cos [id D]\n",
      "    └─ Mul [id E]\n",
      "       ├─ ExpandDims{axis=0} [id F]\n",
      "       │  └─ 6.283185307179586 [id G]\n",
      "       └─ b [id H]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x105c9e0e0>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from pytensor.graph.replace import graph_replace\n",
    "\n",
    "\n",
    "z2 = graph_replace(d, {c: pt.cos(b)})\n",
    "z2.dprint()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ceffb26",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "This might seem like a silly example, but one-to-one transformations are quite common! For example, suppose we had a graph with a random variable, and we wanted to replace the random variable with a root variable, transforming it into a function. We can do that with graph replace too:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "5856f780",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Exp [id A]\n",
      " └─ normal_rv{\"(),()->()\"}.1 [id B]\n",
      "    ├─ RNG(<Generator(PCG64) at 0x12C30BAE0>) [id C]\n",
      "    ├─ [10] [id D]\n",
      "    ├─ ExpandDims{axis=0} [id E]\n",
      "    │  └─ 0 [id F]\n",
      "    └─ ExpandDims{axis=0} [id G]\n",
      "       └─ 1 [id H]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x105c9e0e0>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rng = pt.random.rng(\"rng\")\n",
    "next_rng, x = rng.normal(loc=0, scale=1, size=(10,))\n",
    "y = pt.exp(x)\n",
    "y.dprint()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "0271d4bd",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Exp [id A]\n",
      " └─ x_input [id B]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x105c9e0e0>"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x_input = pt.tensor(\"x_input\", shape=(10,))\n",
    "y2 = graph_replace(y, {x: x_input})\n",
    "y2.dprint()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08ccb932",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "#### Vectorization\n",
    "\n",
    "A less trivial type of graph replacement is *vectorization*. This is also a one-to-one replacement, but this time things are somewhat more complicated, because it changes the meaning of the graph in a very structured manner. Pytensor will automatically reason about the broadcasting operations that need to happen to accomplish the replacement.\n",
    "\n",
    "This time, I make a graph where all the inputs are specifically declared to be scalars. Notice how this information flows down the graph: since `x` is scalar, and the `Op`s `log` and `add` are elemwise, the output `y` is also a scalar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "05efd44a",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Log [id A] <Scalar(float64, shape=())>\n",
      " └─ Add [id B] <Scalar(float64, shape=())>\n",
      "    ├─ 1 [id C] <Scalar(int8, shape=())>\n",
      "    └─ x [id D] <Scalar(float64, shape=())>\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x105c9e0e0>"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = pt.dscalar(\"x\")\n",
    "y = pt.log(1 + x)\n",
    "y.dprint(print_type=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd514d05",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Suppose that instead of a scalar `x`, we want to use this function for a vector input. If we naively try to `graph_replace` here, we will get a shape error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "2af1c20e",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cannot convert Type Vector(float64, shape=(?,)) (of Variable x_vec) into Type Scalar(float64, shape=()). You can try to manually convert x_vec into a Scalar(float64, shape=()).\n"
     ]
    }
   ],
   "source": [
    "x_vec = pt.tensor(\"x_vec\", shape=(None,))\n",
    "try:\n",
    "    graph_replace(y, {x: x_vec})\n",
    "except Exception as e:\n",
    "    print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bb040e9",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Instead, use `vectorize_graph`!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "429d224f",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Log [id A] <Vector(float64, shape=(?,))>\n",
      " └─ Add [id B] <Vector(float64, shape=(?,))>\n",
      "    ├─ ExpandDims{axis=0} [id C] <Vector(int8, shape=(1,))>\n",
      "    │  └─ 1 [id D] <Scalar(int8, shape=())>\n",
      "    └─ x_vec [id E] <Vector(float64, shape=(?,))>\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x105c9e0e0>"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from pytensor.graph.replace import vectorize_graph\n",
    "\n",
    "\n",
    "y_vec = vectorize_graph(y, {x: x_vec})\n",
    "y_vec.dprint(print_type=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "539a6d85",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Rewrites\n",
    "\n",
    "Rewrites are at the heart of pytensor's usefulness. Pytensor maintains large databases of useful graph transformations. These can be applied to acheive a number of goals."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff07ef2c",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "#### Scalarization\n",
    "\n",
    "We can even undo vectorization -- scalarization?\n",
    "\n",
    "Here's the original graph:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "82e05376",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Subtensor{i} [id A]\n",
      " ├─ Log [id B]\n",
      " │  └─ Add [id C]\n",
      " │     ├─ ExpandDims{axis=0} [id D]\n",
      " │     │  └─ 1 [id E]\n",
      " │     └─ x_vec [id F]\n",
      " └─ 0 [id G]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x105c9e0e0>"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y_vec[0].dprint()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a5bab67",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Notice how the indexing operation has been pushed down to the input!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "e3c49b4c",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Log [id A]\n",
      " └─ Add [id B]\n",
      "    ├─ 1.0 [id C]\n",
      "    └─ Subtensor{i} [id D]\n",
      "       ├─ x_vec [id E]\n",
      "       └─ 0 [id F]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x105c9e0e0>"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from pytensor.graph import rewrite_graph\n",
    "\n",
    "\n",
    "rewrite_graph(y_vec[0]).dprint()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6dce4025",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "#### Stabilization\n",
    "\n",
    "Calling `rewrite_graph` without any arguments will trigger many rewrites to be applied. If you have a more specific objective, you can ask for specific types of rewrites. For example, the `stabilize` tag includes rewrites that can transform your graph into a form that is more numerically stable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "c1b3d7c7",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Log [id A] <Vector(float64, shape=(?,))>\n",
      " └─ Add [id B] <Vector(float64, shape=(?,))>\n",
      "    ├─ ExpandDims{axis=0} [id C] <Vector(int8, shape=(1,))>\n",
      "    │  └─ 1 [id D] <Scalar(int8, shape=())>\n",
      "    └─ x [id E] <Vector(float64, shape=(?,))>\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x105c9e0e0>"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = pt.tensor(\"x\", shape=(None,))\n",
    "y = pt.log(1 + x)\n",
    "y.dprint(print_type=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "91174dfd",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Log1p [id A]\n",
      " └─ x [id B]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x105c9e0e0>"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "stable_y = rewrite_graph(y, include=(\"stabilize\",))\n",
    "stable_y.dprint()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e65bb491",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "source": [
    "There are many more types of rewrites, and we will see them shortly. But first we need to do a Quinten Tarintino and show things somewhat out of order. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "290f088c",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Graph-to-graph transformations\n",
    "\n",
    "The most powerful feature of pytensor is to take a graph and return another graph. \n",
    "\n",
    "We've already seen one example of this in `vectorization`. We included it above because it \"looks like\" `graph_replace`, in the sense that you target a single node for replacement. But as we saw, `vectorize_graph` returns an entirely new graph, with new shapes.\n",
    "\n",
    "The canonical example of a graph-to-graph transformation is automatic differentiation. If we know the derivative of every `Op` in our graph, we work backwards and follow the chain rule to construct a gradient graph from the graph of a scalar loss function.\n",
    "\n",
    "The cryptic Op called `Second` means: keep the second input after broadcasting the shape with the first.\n",
    "It's the same as `np.broadcast_arrays(x, y)[1]`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "c455a4c9",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True_div [id A]\n",
      " ├─ Second [id B]\n",
      " │  ├─ Log1p [id C]\n",
      " │  │  └─ x [id D]\n",
      " │  └─ ExpandDims{axis=0} [id E]\n",
      " │     └─ Second [id F]\n",
      " │        ├─ Sum{axes=None} [id G]\n",
      " │        │  └─ Log1p [id C]\n",
      " │        │     └─ ···\n",
      " │        └─ 1.0 [id H]\n",
      " └─ Add [id I]\n",
      "    ├─ ExpandDims{axis=0} [id J]\n",
      "    │  └─ 1 [id K]\n",
      "    └─ x [id D]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x105c9e0e0>"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from pytensor.gradient import grad\n",
    "\n",
    "\n",
    "grad_y = grad(stable_y.sum(), wrt=x)\n",
    "grad_y.dprint()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c0484ee",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "#### Quick digression: back to rewrites!\n",
    "\n",
    "Gradient graphs tend to be complex, so it's nice to simplify them. \n",
    "\n",
    "One type of simplification is *canonicalization*. It converts a graph into a \"standard\" form. Other rewrites can expect and reason from this form. I think you'll agree this form is much nicer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "1bef9b01",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True_div [id A]\n",
      " ├─ [1.] [id B]\n",
      " └─ Add [id C]\n",
      "    ├─ [1.] [id B]\n",
      "    └─ x [id D]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x105c9e0e0>"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rewrite_graph(grad_y, include=(\"canonicalize\",)).dprint()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97a7c16b",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "We can also compose rewrites! In this next example, we both canonicalize, then apply specializations to get computational speedups.\n",
    "\n",
    "This also reveals the \"final form\" of the gradient: $\\log(1 + x) = \\frac{1}{1 + x}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "0262e4ba",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reciprocal [id A]\n",
      " └─ Add [id B]\n",
      "    ├─ [1.] [id C]\n",
      "    └─ x [id D]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x105c9e0e0>"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rewrite_graph(grad_y, include=(\"canonicalize\", \"specialize\")).dprint()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d492543e",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "#### Back to transformations -- logp inference\n",
    "\n",
    "Another important graph-to-graph transformation is *automatic logp inference*. PyMC knows how to transform a generatve graph (forward draws) into a logp graph (backwards inference). Let's look at a simple example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "37ed37be",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "normal_rv{\"(),()->()\"}.1 [id A] 'z'\n",
      " ├─ RNG(<Generator(PCG64) at 0x136F5DC40>) [id B]\n",
      " ├─ NoneConst{None} [id C]\n",
      " ├─ [0 0] [id D]\n",
      " └─ [1 2] [id E]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x105c9e0e0>"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pymc as pm\n",
    "\n",
    "\n",
    "with pm.Model() as model:\n",
    "    z = pm.Normal(name=\"z\", mu=np.array([0, 0]), sigma=np.array([1, 2]))\n",
    "\n",
    "pytensor.dprint(z)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43b8a6da",
   "metadata": {},
   "source": [
    "We can now use  {func}`~pymc.logp` to compute the log-probability of a value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "8e3194d8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Check{sigma > 0} [id A] 'z_logprob'\n",
      " ├─ Sub [id B]\n",
      " │  ├─ Sub [id C]\n",
      " │  │  ├─ Mul [id D]\n",
      " │  │  │  ├─ ExpandDims{axis=0} [id E]\n",
      " │  │  │  │  └─ -0.5 [id F]\n",
      " │  │  │  └─ Pow [id G]\n",
      " │  │  │     ├─ True_div [id H]\n",
      " │  │  │     │  ├─ Sub [id I]\n",
      " │  │  │     │  │  ├─ z_value [id J]\n",
      " │  │  │     │  │  └─ [0 0] [id K]\n",
      " │  │  │     │  └─ [1 2] [id L]\n",
      " │  │  │     └─ ExpandDims{axis=0} [id M]\n",
      " │  │  │        └─ 2 [id N]\n",
      " │  │  └─ ExpandDims{axis=0} [id O]\n",
      " │  │     └─ Log [id P]\n",
      " │  │        └─ Sqrt [id Q]\n",
      " │  │           └─ 6.283185307179586 [id R]\n",
      " │  └─ Log [id S]\n",
      " │     └─ [1 2] [id L]\n",
      " └─ All{axes=None} [id T]\n",
      "    └─ MakeVector{dtype='bool'} [id U]\n",
      "       └─ All{axes=None} [id V]\n",
      "          └─ Gt [id W]\n",
      "             ├─ [1 2] [id L]\n",
      "             └─ ExpandDims{axis=0} [id X]\n",
      "                └─ 0 [id Y]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x105c9e0e0>"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z_value = pt.vector(name=\"z_value\")\n",
    "z_logp = pm.logp(rv=z, value=z_value)\n",
    "\n",
    "pytensor.dprint(z_logp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "942aadb1",
   "metadata": {},
   "source": [
    "Observe that we still get a graph, we do not do any computation yet. To do a evaluation on a given value, we can use the `eval` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "1a0f1ff3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Check{sigma > 0}(Sub.0, All{axes=None}.0)"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z_logp.owner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "a6aff350",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([-0.91893853, -1.61208571])"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z_logp.eval({z_value: [0, 0]})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43339908",
   "metadata": {},
   "source": [
    "To \"verify\" the computation, we compute the same quantity using `scipy.stats.norm.logpdf`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "3453b57f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([-0.91893853, -1.61208571])"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import scipy\n",
    "\n",
    "\n",
    "scipy.stats.norm.logpdf(\n",
    "    x=np.array([0, 0]), loc=np.array([0, 0]), scale=np.array([1, 2])\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65c3537c",
   "metadata": {},
   "source": [
    "The values match as expected."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e888b20",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Composability\n",
    "\n",
    "It is important to note that all of these transformation operations are *composable*. We already saw an example of this when we applied graph rewrites to a gradient graph. We can also vectorize gradients, or take the gradient of a vectorize graph. Or we can replace subgraphs before or after applying rewrites. Because all of these operations take in a graph and return a graph, we're always able to chain together graph operations. \n",
    "\n",
    "Sometimes, you'll be surprised by what you end up with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "f44aed85",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Log [id A]\n",
      " └─ Add [id B]\n",
      "    ├─ ExpandDims{axis=0} [id C]\n",
      "    │  └─ 1 [id D]\n",
      "    └─ x [id E]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x105c9e0e0>"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = pt.tensor(\"x\", shape=(3,))\n",
    "y = pt.log(1 + x)\n",
    "y.dprint()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa1eb391",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "The graph we've been working with is $y = \\log(1 + x)$. Let's transform it to be $y = \\log(1 + \\exp(x))$, then rewrite the expression for stability.\n",
    "\n",
    "First replace $x$ by $\\exp(x)$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "5dc415eb",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Log [id A]\n",
      " └─ Add [id B]\n",
      "    ├─ ExpandDims{axis=0} [id C]\n",
      "    │  └─ 1 [id D]\n",
      "    └─ Exp [id E]\n",
      "       └─ x [id F]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x105c9e0e0>"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_y = graph_replace(y, replace={x: pt.exp(x)})\n",
    "new_y.dprint()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d418c581",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Then optimize the graph for stability:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "94ca84f9",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Scalar_softplus [id A]\n",
      " └─ x [id B]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x105c9e0e0>"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rewrite_graph(new_y, include=(\"stabilize\",)).dprint()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97e2650b",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Did you guess the result would be just a single Op?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93ed3570-45f3-4df4-bbfd-4ffa0497d30e",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## (3) Compilation\n",
    "\n",
    "All this is fun and dandy but only useful if we actually use it compute stuff! \n",
    "\n",
    "PyTensor provides a critical non-composable graph operation: `function`, which converts a pytensor graph into a callable python object that takes concrete inputs and returns concrete outputs. \n",
    "\n",
    "By default it runs an extensive database of rewrites to try and optimize the computational graph, and then compiles to C (technically a mix of C and Python if not all operations have a C implementation). See https://pytensor.readthedocs.io/en/latest/extending/pipeline.html for a bit more detail.\n",
    "\n",
    "As with anything remotely useful in Python, when it comes to work you want to [STAY OUT OF PYTHON](https://www.youtube.com/watch?v=vVUnCXKuNOg) as much as possible."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "7d74a4d9-2079-4462-81a3-a4c17a271f29",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cos [id A]\n",
      " └─ Squeeze{axes=[0, 1]} [id B]\n",
      "    └─ Blockwise{dot, (m,k),(k,n)->(m,n)} [id C]\n",
      "       ├─ ExpandDims{axis=0} [id D]\n",
      "       │  └─ Exp [id E]\n",
      "       │     └─ Sin [id F]\n",
      "       │        └─ x [id G]\n",
      "       └─ ExpandDims{axis=1} [id H]\n",
      "          └─ Exp [id E]\n",
      "             └─ ···\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x105c9e0e0>"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = pt.vector(\"x\", shape=(None,))\n",
    "z = pt.exp(pt.sin(x))\n",
    "out = pt.cos((z[None, :] @ z[:, None]).squeeze())\n",
    "out.dprint()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9446b3e",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "As noted, use `pytensor.function` to compile a graph into an executable program. You need to pass a **list** of inputs (even if it has no inputs, you still need an empty list!), and outputs (this *can* be a list, but isn't required to be."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "c98b61df45d6a2bb",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "y_fn = pytensor.function([x], out)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e2727ea",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "What we get back is a `Function` object. This is a wrapper around a call out to a compiled `C` program that now lives in some cache folder somewhere on your computer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "5a8f3eca-51bf-4a3d-a656-85df87135c34",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "pytensor.compile.function.types.Function"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(y_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8f7e69c",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "And now we can use it like any other python function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "90f520a5-eec6-42ed-abad-26204cac5868",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array(0.40678993)"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y_fn(np.random.randn(3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "3badd33e-ddb3-4aca-a27f-49275b1b994b",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array(0.03215766)"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y_fn(np.random.randn(5))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7574ded",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "We can also inspect it to see the final graph that was compiled, after applying rewrites. Notice that `dot` became `CGemv`, so we're using the correct BLAS routine for the inputs provided. It also did some *loop fusion* by compiling a composite inner graph. That is, if we give an array input, rather than looping over it once to compute `exp` of each element, then looping over it again to compute `sin` of each element, we instead loop only once, and compute $\\exp \\circ \\sin$ of each element."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "e300df411e0aff49",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cos [id A] d={0: [0]} 5\n",
      " └─ Squeeze{axis=0} [id B] 4\n",
      "    └─ CGemv{inplace} [id C] d={0: [0]} 3\n",
      "       ├─ AllocEmpty{dtype='float64'} [id D] 1\n",
      "       │  └─ 1 [id E]\n",
      "       ├─ 1.0 [id F]\n",
      "       ├─ ExpandDims{axis=0} [id G] 2\n",
      "       │  └─ Composite{exp(sin(i0))} [id H] 0\n",
      "       │     └─ x [id I]\n",
      "       ├─ Composite{exp(sin(i0))} [id H] 0\n",
      "       │  └─ ···\n",
      "       └─ 0.0 [id J]\n",
      "\n",
      "Inner graphs:\n",
      "\n",
      "Composite{exp(sin(i0))} [id H]\n",
      " ← exp [id K] 'o0'\n",
      "    └─ sin [id L]\n",
      "       └─ i0 [id M]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x105c9e0e0>"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y_fn.dprint(print_destroy_map=True)  # Some memory aliasing optimizations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b4adffa",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "By passing `print_destroy_map=True` to `dprint`, we also get to see where memory buffers are being reused. This shows up in the graph with the line `d={x: [y]}`, where \"x\" is an output, and \"y\" is an input that will be re-used as a buffer. We can see this in two places:\n",
    "\n",
    "- In the `CGemv` line, we have `d={0: [0]}`. This means that the first output is being allocated to the memory used for the first input. In this case, we save allocation of a length 1 array, which isn't so impressive.\n",
    "\n",
    "- More interestingly, we see it at the top line, `Cos [id A] d = {0: [0]}`. This means that we are doing the `Cos` directly on the squeezed output of the CGemv operation! This is equivalent to `np.cos(x, out=x)`, but the user doesn't need to remember to do it themselves, and the code stays more readable "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5936a640-43ca-4a37-85e5-66feeda4558f",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "PyTensor can also delegate compilation to other libraries in town, namely Numba, JAX, and PyTorch (latter still under active development). \n",
    "\n",
    "Notice that this is another kind of graph-to-\"graph\" transformation, that we can only do because we have access to the whole static computation. Every individual `Op` knows what it means to become a Numba program, so we can just walk across the graph and generate the appropriate code.\n",
    "\n",
    "Sometimes, this means that certain rewrites aren't applied! In this case, Numba does it's own BLAS optimizations. So we don't rewrite `dot -> CGemv`, we just leave it as dot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "719fdfbe-790d-41c5-9d4a-398697f01926",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cos [id A] d={0: [0]} 5\n",
      " └─ Squeeze{axes=[0, 1]} [id B] 4\n",
      "    └─ dot [id C] 3\n",
      "       ├─ ExpandDims{axis=0} [id D] 2\n",
      "       │  └─ Composite{exp(sin(i0))} [id E] 0\n",
      "       │     └─ x [id F]\n",
      "       └─ ExpandDims{axis=1} [id G] 1\n",
      "          └─ Composite{exp(sin(i0))} [id E] 0\n",
      "             └─ ···\n",
      "\n",
      "Inner graphs:\n",
      "\n",
      "Composite{exp(sin(i0))} [id E]\n",
      " ← exp [id H] 'o0'\n",
      "    └─ sin [id I]\n",
      "       └─ i0 [id J]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x105c9e0e0>"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y_numba_fn = pytensor.function([x], out, mode=\"NUMBA\")\n",
    "y_numba_fn.dprint(print_destroy_map=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "f9fcb852-49b3-4442-86e8-eccab65f3834",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 96.3 ms, sys: 7.66 ms, total: 104 ms\n",
      "Wall time: 107 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array(-0.82110244)"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "y_numba_fn(\n",
    "    np.random.randn(3)\n",
    ")  # first time takes long, jit compilation actually happening"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "b007790d-fa14-4ac4-823e-94cf2df15e35",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 141 μs, sys: 75 μs, total: 216 μs\n",
      "Wall time: 217 μs\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array(-0.54402517)"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "y_numba_fn(np.random.randn(5))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eab1d41c",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## A Full Example: Logistic Regression with Gradient Descent\n",
    "\n",
    "To show everything in action together, let's look at how we would write a binary classification model, how we can train it using gradient descent, and how we can use the final program we get back.\n",
    "\n",
    "First, let's set up the symbolic inputs. These will be the input data $X$ and targets $y$, as well as the initial values for alpha and beta. Our model will be:\n",
    "\n",
    "$$ p = \\sigma^{-1}(\\alpha + X \\beta) $$\n",
    "\n",
    "Where $\\sigma$ is the logistic function.\n",
    "\n",
    "And we'll choose parameters $\\alpha$ and $\\beta$ to minimize the binary cross entropy between $p$ and the target labels, which will just be:\n",
    "\n",
    "$$\\mathcal{L} = -\\frac{1}{N} \\sum_{i=0}^N y_i \\log(p_i) + (1 - y_i) \\log(1 - p_i)$$\n",
    "\n",
    "If there are any Bayesians in the audience, you will also recognize this as the negative log-likelihood of a Bernoulli GLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "e448cd31",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "X_pt = pt.tensor(\"X\", shape=(None, None))\n",
    "y_pt = pt.tensor(\"y\", shape=(None,))\n",
    "alpha_pt = pt.tensor(\"alpha\", shape=())\n",
    "beta_pt = pt.tensor(\"beta\", shape=(None,))\n",
    "\n",
    "p = pt.sigmoid(alpha_pt + X_pt @ beta_pt)\n",
    "p.name = \"p_class_0\"\n",
    "\n",
    "loss = -(y_pt * pt.log(p) + (1 - y_pt) * pt.log(1 - p)).mean()\n",
    "loss.name = \"cross_entropy\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb622fa5",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "To convince you that the text version of dprint is nicer, here's the big graphic plot for our loss function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "d52561b6",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<svg xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\" width=\"901pt\" height=\"1522pt\" viewBox=\"0.00 0.00 901.00 1522.00\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 1517.5)\">\n",
       "<title>G</title>\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-1517.5 896.83,-1517.5 896.83,4 -4,4\"/>\n",
       "<!-- 5226949680 -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>5226949680</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"113.75\" cy=\"-1422.5\" rx=\"93.83\" ry=\"18\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"113.75\" y=\"-1417.45\" font-family=\"Times,serif\" font-size=\"14.00\">ExpandDims{axis=1}</text>\n",
       "</g>\n",
       "<!-- 4340722064 -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>4340722064</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"226.75\" cy=\"-1334\" rx=\"147.06\" ry=\"18\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"226.75\" y=\"-1328.95\" font-family=\"Times,serif\" font-size=\"14.00\">Blockwise{dot, (m,k),(k,n)-&gt;(m,n)}</text>\n",
       "</g>\n",
       "<!-- 5226949680&#45;&gt;4340722064 -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>5226949680-&gt;4340722064</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M115.2,-1404.23C116.99,-1393.16 120.95,-1379.2 129.75,-1370 135.28,-1364.22 141.83,-1359.39 148.89,-1355.35\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"150.09,-1358.67 157.41,-1351.02 146.91,-1352.43 150.09,-1358.67\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"215.25\" y=\"-1373.2\" font-family=\"Times,serif\" font-size=\"14.00\">1 Matrix(float64, shape=(?, 1))</text>\n",
       "</g>\n",
       "<!-- 5036693584 -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>5036693584</title>\n",
       "<polygon fill=\"green\" stroke=\"black\" points=\"227.5,-1513.5 0,-1513.5 0,-1477.5 227.5,-1477.5 227.5,-1513.5\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"113.75\" y=\"-1490.45\" font-family=\"Times,serif\" font-size=\"14.00\">name=beta Vector(float64, shape=(?,))</text>\n",
       "</g>\n",
       "<!-- 5036693584&#45;&gt;5226949680 -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>5036693584-&gt;5226949680</title>\n",
       "<path fill=\"none\" stroke=\"blue\" d=\"M113.75,-1477.31C113.75,-1469.73 113.75,-1460.6 113.75,-1452.04\"/>\n",
       "<polygon fill=\"blue\" stroke=\"blue\" points=\"117.25,-1452.04 113.75,-1442.04 110.25,-1452.04 117.25,-1452.04\"/>\n",
       "</g>\n",
       "<!-- 4340721632 -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>4340721632</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"282.75\" cy=\"-1245.5\" rx=\"75.92\" ry=\"18\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"282.75\" y=\"-1240.45\" font-family=\"Times,serif\" font-size=\"14.00\">Squeeze{axis=1}</text>\n",
       "</g>\n",
       "<!-- 4340722064&#45;&gt;4340721632 -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>4340722064-&gt;4340721632</title>\n",
       "<path fill=\"none\" stroke=\"blue\" d=\"M237.81,-1315.91C245.71,-1303.71 256.5,-1287.05 265.52,-1273.12\"/>\n",
       "<polygon fill=\"blue\" stroke=\"blue\" points=\"268.39,-1275.12 270.89,-1264.82 262.51,-1271.32 268.39,-1275.12\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"340.02\" y=\"-1284.7\" font-family=\"Times,serif\" font-size=\"14.00\">Matrix(float64, shape=(?, 1))</text>\n",
       "</g>\n",
       "<!-- 5224545120 -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>5224545120</title>\n",
       "<polygon fill=\"green\" stroke=\"black\" points=\"453.5,-1440.5 226,-1440.5 226,-1404.5 453.5,-1404.5 453.5,-1440.5\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"339.75\" y=\"-1417.45\" font-family=\"Times,serif\" font-size=\"14.00\">name=X Matrix(float64, shape=(?, ?))</text>\n",
       "</g>\n",
       "<!-- 5224545120&#45;&gt;4340722064 -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>5224545120-&gt;4340722064</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M331.31,-1404.39C325.19,-1393.39 316.02,-1379.45 304.75,-1370 298.65,-1364.89 291.69,-1360.38 284.48,-1356.45\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"286.43,-1353.51 275.92,-1352.13 283.27,-1359.76 286.43,-1353.51\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"323.05\" y=\"-1373.2\" font-family=\"Times,serif\" font-size=\"14.00\">0</text>\n",
       "</g>\n",
       "<!-- 4340613344 -->\n",
       "<g id=\"node8\" class=\"node\">\n",
       "<title>4340613344</title>\n",
       "<ellipse fill=\"#ffaabb\" stroke=\"black\" cx=\"375.75\" cy=\"-1157\" rx=\"27\" ry=\"18\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"375.75\" y=\"-1151.95\" font-family=\"Times,serif\" font-size=\"14.00\">Add</text>\n",
       "</g>\n",
       "<!-- 4340721632&#45;&gt;4340613344 -->\n",
       "<g id=\"edge7\" class=\"edge\">\n",
       "<title>4340721632-&gt;4340613344</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M280.75,-1227.08C280.45,-1216.23 281.88,-1202.57 289.25,-1193 301,-1177.73 320.51,-1169.09 338.02,-1164.22\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"338.78,-1167.64 347.67,-1161.89 337.13,-1160.84 338.78,-1167.64\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"368\" y=\"-1196.2\" font-family=\"Times,serif\" font-size=\"14.00\">1 Vector(float64, shape=(?,))</text>\n",
       "</g>\n",
       "<!-- 4340720816 -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>4340720816</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"492.75\" cy=\"-1245.5\" rx=\"93.83\" ry=\"18\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"492.75\" y=\"-1240.45\" font-family=\"Times,serif\" font-size=\"14.00\">ExpandDims{axis=0}</text>\n",
       "</g>\n",
       "<!-- 4340720816&#45;&gt;4340613344 -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>4340720816-&gt;4340613344</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M480.08,-1227.18C471.54,-1216.37 459.54,-1202.72 446.75,-1193 435.41,-1184.38 421.62,-1177.01 409.21,-1171.28\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"410.71,-1168.11 400.15,-1167.29 407.89,-1174.52 410.71,-1168.11\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"542.71\" y=\"-1196.2\" font-family=\"Times,serif\" font-size=\"14.00\">0 Vector(float64, shape=(1,))</text>\n",
       "</g>\n",
       "<!-- 5036697616 -->\n",
       "<g id=\"node7\" class=\"node\">\n",
       "<title>5036697616</title>\n",
       "<polygon fill=\"green\" stroke=\"black\" points=\"615.62,-1352 391.88,-1352 391.88,-1316 615.62,-1316 615.62,-1352\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"503.75\" y=\"-1328.95\" font-family=\"Times,serif\" font-size=\"14.00\">name=alpha Scalar(float64, shape=())</text>\n",
       "</g>\n",
       "<!-- 5036697616&#45;&gt;4340720816 -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>5036697616-&gt;4340720816</title>\n",
       "<path fill=\"none\" stroke=\"blue\" d=\"M501.58,-1315.91C500.09,-1304.26 498.1,-1288.55 496.38,-1275.02\"/>\n",
       "<polygon fill=\"blue\" stroke=\"blue\" points=\"499.88,-1274.83 495.15,-1265.35 492.94,-1275.71 499.88,-1274.83\"/>\n",
       "</g>\n",
       "<!-- 4340722208 -->\n",
       "<g id=\"node9\" class=\"node\">\n",
       "<title>4340722208</title>\n",
       "<ellipse fill=\"#ffaabb\" stroke=\"black\" cx=\"375.75\" cy=\"-1068.5\" rx=\"42.65\" ry=\"18\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"375.75\" y=\"-1063.45\" font-family=\"Times,serif\" font-size=\"14.00\">Sigmoid</text>\n",
       "</g>\n",
       "<!-- 4340613344&#45;&gt;4340722208 -->\n",
       "<g id=\"edge8\" class=\"edge\">\n",
       "<title>4340613344-&gt;4340722208</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M375.75,-1138.91C375.75,-1127.26 375.75,-1111.55 375.75,-1098.02\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"379.25,-1098.36 375.75,-1088.36 372.25,-1098.36 379.25,-1098.36\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"449.62\" y=\"-1107.7\" font-family=\"Times,serif\" font-size=\"14.00\">Vector(float64, shape=(?,))</text>\n",
       "</g>\n",
       "<!-- 4340462464 -->\n",
       "<g id=\"node10\" class=\"node\">\n",
       "<title>4340462464</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"504.88,-998 246.62,-998 246.62,-962 504.88,-962 504.88,-998\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"375.75\" y=\"-974.95\" font-family=\"Times,serif\" font-size=\"14.00\">name=p_class_0 Vector(float64, shape=(?,))</text>\n",
       "</g>\n",
       "<!-- 4340722208&#45;&gt;4340462464 -->\n",
       "<g id=\"edge9\" class=\"edge\">\n",
       "<title>4340722208-&gt;4340462464</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M375.75,-1050.41C375.75,-1038.76 375.75,-1023.05 375.75,-1009.52\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"379.25,-1009.86 375.75,-999.86 372.25,-1009.86 379.25,-1009.86\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"449.62\" y=\"-1019.2\" font-family=\"Times,serif\" font-size=\"14.00\">Vector(float64, shape=(?,))</text>\n",
       "</g>\n",
       "<!-- 4340462992 -->\n",
       "<g id=\"node13\" class=\"node\">\n",
       "<title>4340462992</title>\n",
       "<ellipse fill=\"#ffaabb\" stroke=\"black\" cx=\"235.75\" cy=\"-891.5\" rx=\"27\" ry=\"18\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"235.75\" y=\"-886.45\" font-family=\"Times,serif\" font-size=\"14.00\">Sub</text>\n",
       "</g>\n",
       "<!-- 4340462464&#45;&gt;4340462992 -->\n",
       "<g id=\"edge12\" class=\"edge\">\n",
       "<title>4340462464-&gt;4340462992</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M364.65,-961.94C356.51,-950.67 344.51,-936.4 330.75,-927.5 307.73,-912.6 297.32,-919.38 271.75,-909.5 270.22,-908.91 268.66,-908.28 267.1,-907.63\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"268.79,-904.55 258.23,-903.7 265.96,-910.95 268.79,-904.55\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"352.57\" y=\"-930.7\" font-family=\"Times,serif\" font-size=\"14.00\">1</text>\n",
       "</g>\n",
       "<!-- 5228886592 -->\n",
       "<g id=\"node20\" class=\"node\">\n",
       "<title>5228886592</title>\n",
       "<ellipse fill=\"#ffaabb\" stroke=\"black\" cx=\"806.75\" cy=\"-891.5\" rx=\"50.33\" ry=\"18\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"806.75\" y=\"-886.45\" font-family=\"Times,serif\" font-size=\"14.00\">Log id=12</text>\n",
       "</g>\n",
       "<!-- 4340462464&#45;&gt;5228886592 -->\n",
       "<g id=\"edge19\" class=\"edge\">\n",
       "<title>4340462464-&gt;5228886592</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M484.17,-961.57C558.69,-948.9 659.68,-930.36 747.75,-909.5 751.12,-908.7 754.58,-907.83 758.05,-906.93\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"758.93,-910.32 767.66,-904.32 757.09,-903.56 758.93,-910.32\"/>\n",
       "</g>\n",
       "<!-- 4340465008 -->\n",
       "<g id=\"node11\" class=\"node\">\n",
       "<title>4340465008</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"114.75\" cy=\"-980\" rx=\"113.79\" ry=\"18\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"114.75\" y=\"-974.95\" font-family=\"Times,serif\" font-size=\"14.00\">ExpandDims{axis=0} id=6</text>\n",
       "</g>\n",
       "<!-- 4340465008&#45;&gt;4340462992 -->\n",
       "<g id=\"edge11\" class=\"edge\">\n",
       "<title>4340465008-&gt;4340462992</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M138.66,-961.91C158.75,-947.54 187.48,-927.01 208.47,-912\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"210.47,-914.88 216.57,-906.21 206.39,-909.18 210.47,-914.88\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"256.59\" y=\"-930.7\" font-family=\"Times,serif\" font-size=\"14.00\">0 Vector(int8, shape=(1,))</text>\n",
       "</g>\n",
       "<!-- 5228873872 -->\n",
       "<g id=\"node12\" class=\"node\">\n",
       "<title>5228873872</title>\n",
       "<polygon fill=\"green\" stroke=\"black\" points=\"200,-1086.5 29.5,-1086.5 29.5,-1050.5 200,-1050.5 200,-1086.5\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"114.75\" y=\"-1063.45\" font-family=\"Times,serif\" font-size=\"14.00\">val=1 Scalar(int8, shape=())</text>\n",
       "</g>\n",
       "<!-- 5228873872&#45;&gt;4340465008 -->\n",
       "<g id=\"edge10\" class=\"edge\">\n",
       "<title>5228873872-&gt;4340465008</title>\n",
       "<path fill=\"none\" stroke=\"blue\" d=\"M114.75,-1050.41C114.75,-1038.76 114.75,-1023.05 114.75,-1009.52\"/>\n",
       "<polygon fill=\"blue\" stroke=\"blue\" points=\"118.25,-1009.86 114.75,-999.86 111.25,-1009.86 118.25,-1009.86\"/>\n",
       "</g>\n",
       "<!-- 5228875216 -->\n",
       "<g id=\"node14\" class=\"node\">\n",
       "<title>5228875216</title>\n",
       "<ellipse fill=\"#ffaabb\" stroke=\"black\" cx=\"250.75\" cy=\"-803\" rx=\"27\" ry=\"18\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"250.75\" y=\"-797.95\" font-family=\"Times,serif\" font-size=\"14.00\">Log</text>\n",
       "</g>\n",
       "<!-- 4340462992&#45;&gt;5228875216 -->\n",
       "<g id=\"edge13\" class=\"edge\">\n",
       "<title>4340462992-&gt;5228875216</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M238.71,-873.41C240.73,-861.76 243.46,-846.05 245.8,-832.52\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"249.22,-833.29 247.48,-822.84 242.33,-832.1 249.22,-833.29\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"318.47\" y=\"-842.2\" font-family=\"Times,serif\" font-size=\"14.00\">Vector(float64, shape=(?,))</text>\n",
       "</g>\n",
       "<!-- 5228886304 -->\n",
       "<g id=\"node19\" class=\"node\">\n",
       "<title>5228886304</title>\n",
       "<ellipse fill=\"#ffaabb\" stroke=\"black\" cx=\"439.75\" cy=\"-714.5\" rx=\"27\" ry=\"18\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"439.75\" y=\"-709.45\" font-family=\"Times,serif\" font-size=\"14.00\">Mul</text>\n",
       "</g>\n",
       "<!-- 5228875216&#45;&gt;5228886304 -->\n",
       "<g id=\"edge18\" class=\"edge\">\n",
       "<title>5228875216-&gt;5228886304</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M255.6,-785.13C259.76,-773.65 266.93,-759.05 278.25,-750.5 313.83,-723.62 365.74,-716.48 401.03,-715.02\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"400.9,-718.52 410.81,-714.76 400.72,-711.53 400.9,-718.52\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"357\" y=\"-753.7\" font-family=\"Times,serif\" font-size=\"14.00\">1 Vector(float64, shape=(?,))</text>\n",
       "</g>\n",
       "<!-- 5228883184 -->\n",
       "<g id=\"node15\" class=\"node\">\n",
       "<title>5228883184</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"394.75\" cy=\"-891.5\" rx=\"113.79\" ry=\"18\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"394.75\" y=\"-886.45\" font-family=\"Times,serif\" font-size=\"14.00\">ExpandDims{axis=0} id=9</text>\n",
       "</g>\n",
       "<!-- 5228883472 -->\n",
       "<g id=\"node17\" class=\"node\">\n",
       "<title>5228883472</title>\n",
       "<ellipse fill=\"#ffaabb\" stroke=\"black\" cx=\"439.75\" cy=\"-803\" rx=\"49.82\" ry=\"18\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"439.75\" y=\"-797.95\" font-family=\"Times,serif\" font-size=\"14.00\">Sub id=10</text>\n",
       "</g>\n",
       "<!-- 5228883184&#45;&gt;5228883472 -->\n",
       "<g id=\"edge15\" class=\"edge\">\n",
       "<title>5228883184-&gt;5228883472</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M403.64,-873.41C409.93,-861.33 418.48,-844.88 425.69,-831.03\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"428.67,-832.88 430.18,-822.4 422.46,-829.65 428.67,-832.88\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"491.78\" y=\"-842.2\" font-family=\"Times,serif\" font-size=\"14.00\">0 Vector(int8, shape=(1,))</text>\n",
       "</g>\n",
       "<!-- 5228806064 -->\n",
       "<g id=\"node16\" class=\"node\">\n",
       "<title>5228806064</title>\n",
       "<polygon fill=\"green\" stroke=\"black\" points=\"693,-998 522.5,-998 522.5,-962 693,-962 693,-998\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"607.75\" y=\"-974.95\" font-family=\"Times,serif\" font-size=\"14.00\">val=1 Scalar(int8, shape=())</text>\n",
       "</g>\n",
       "<!-- 5228806064&#45;&gt;5228883184 -->\n",
       "<g id=\"edge14\" class=\"edge\">\n",
       "<title>5228806064-&gt;5228883184</title>\n",
       "<path fill=\"none\" stroke=\"blue\" d=\"M564.9,-961.6C530.29,-947.55 481.52,-927.74 445.06,-912.93\"/>\n",
       "<polygon fill=\"blue\" stroke=\"blue\" points=\"446.63,-909.79 436.05,-909.27 444,-916.28 446.63,-909.79\"/>\n",
       "</g>\n",
       "<!-- 5228883472&#45;&gt;5228886304 -->\n",
       "<g id=\"edge17\" class=\"edge\">\n",
       "<title>5228883472-&gt;5228886304</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M439.75,-784.91C439.75,-773.26 439.75,-757.55 439.75,-744.02\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"443.25,-744.36 439.75,-734.36 436.25,-744.36 443.25,-744.36\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"518.5\" y=\"-753.7\" font-family=\"Times,serif\" font-size=\"14.00\">0 Vector(float64, shape=(?,))</text>\n",
       "</g>\n",
       "<!-- 5224661216 -->\n",
       "<g id=\"node18\" class=\"node\">\n",
       "<title>5224661216</title>\n",
       "<polygon fill=\"green\" stroke=\"black\" points=\"738.62,-909.5 526.88,-909.5 526.88,-873.5 738.62,-873.5 738.62,-909.5\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"632.75\" y=\"-886.45\" font-family=\"Times,serif\" font-size=\"14.00\">name=y Vector(float64, shape=(?,))</text>\n",
       "</g>\n",
       "<!-- 5224661216&#45;&gt;5228883472 -->\n",
       "<g id=\"edge16\" class=\"edge\">\n",
       "<title>5224661216-&gt;5228883472</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M614.48,-873.01C601.93,-861.86 584.41,-847.88 566.75,-839 544.36,-827.75 517.87,-819.71 494.89,-814.21\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"495.93,-810.85 485.4,-812.04 494.37,-817.68 495.93,-810.85\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"595.39\" y=\"-842.2\" font-family=\"Times,serif\" font-size=\"14.00\">1</text>\n",
       "</g>\n",
       "<!-- 5228886880 -->\n",
       "<g id=\"node21\" class=\"node\">\n",
       "<title>5228886880</title>\n",
       "<ellipse fill=\"#ffaabb\" stroke=\"black\" cx=\"632.75\" cy=\"-803\" rx=\"51.35\" ry=\"18\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"632.75\" y=\"-797.95\" font-family=\"Times,serif\" font-size=\"14.00\">Mul id=13</text>\n",
       "</g>\n",
       "<!-- 5224661216&#45;&gt;5228886880 -->\n",
       "<g id=\"edge20\" class=\"edge\">\n",
       "<title>5224661216-&gt;5228886880</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M632.75,-873.41C632.75,-861.76 632.75,-846.05 632.75,-832.52\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"636.25,-832.86 632.75,-822.86 629.25,-832.86 636.25,-832.86\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"636.12\" y=\"-842.2\" font-family=\"Times,serif\" font-size=\"14.00\">0</text>\n",
       "</g>\n",
       "<!-- 5228887168 -->\n",
       "<g id=\"node22\" class=\"node\">\n",
       "<title>5228887168</title>\n",
       "<ellipse fill=\"#ffaabb\" stroke=\"black\" cx=\"529.75\" cy=\"-626\" rx=\"51.35\" ry=\"18\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"529.75\" y=\"-620.95\" font-family=\"Times,serif\" font-size=\"14.00\">Add id=14</text>\n",
       "</g>\n",
       "<!-- 5228886304&#45;&gt;5228887168 -->\n",
       "<g id=\"edge23\" class=\"edge\">\n",
       "<title>5228886304-&gt;5228887168</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M441.31,-696.5C443.1,-685.55 446.95,-671.61 455.25,-662 462.11,-654.05 471.16,-647.79 480.59,-642.89\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"481.91,-646.13 489.5,-638.74 478.96,-639.79 481.91,-646.13\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"534\" y=\"-665.2\" font-family=\"Times,serif\" font-size=\"14.00\">1 Vector(float64, shape=(?,))</text>\n",
       "</g>\n",
       "<!-- 5228886592&#45;&gt;5228886880 -->\n",
       "<g id=\"edge21\" class=\"edge\">\n",
       "<title>5228886592-&gt;5228886880</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M778.41,-876.41C749.3,-861.94 703.79,-839.31 671.39,-823.21\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"673.24,-820.22 662.73,-818.9 670.12,-826.49 673.24,-820.22\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"814.08\" y=\"-842.2\" font-family=\"Times,serif\" font-size=\"14.00\">1 Vector(float64, shape=(?,))</text>\n",
       "</g>\n",
       "<!-- 5228886880&#45;&gt;5228887168 -->\n",
       "<g id=\"edge22\" class=\"edge\">\n",
       "<title>5228886880-&gt;5228887168</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M635.28,-784.9C638.53,-756.3 640.78,-698.36 612.75,-662 605.29,-652.33 594.66,-645.3 583.49,-640.2\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"585.13,-637.09 574.55,-636.61 582.52,-643.59 585.13,-637.09\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"714.94\" y=\"-709.45\" font-family=\"Times,serif\" font-size=\"14.00\">0 Vector(float64, shape=(?,))</text>\n",
       "</g>\n",
       "<!-- 5224665776 -->\n",
       "<g id=\"node23\" class=\"node\">\n",
       "<title>5224665776</title>\n",
       "<ellipse fill=\"brown\" stroke=\"black\" cx=\"418.75\" cy=\"-537.5\" rx=\"33.44\" ry=\"18\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"418.75\" y=\"-532.45\" font-family=\"Times,serif\" font-size=\"14.00\">Shape</text>\n",
       "</g>\n",
       "<!-- 5228887168&#45;&gt;5224665776 -->\n",
       "<g id=\"edge24\" class=\"edge\">\n",
       "<title>5228887168-&gt;5224665776</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M492.41,-613.25C478.69,-607.7 463.67,-600.04 452,-590 444,-583.12 437.2,-573.84 431.88,-565.05\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"435.01,-563.47 427.07,-556.46 428.9,-566.89 435.01,-563.47\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"525.88\" y=\"-576.7\" font-family=\"Times,serif\" font-size=\"14.00\">Vector(float64, shape=(?,))</text>\n",
       "</g>\n",
       "<!-- 5228709840 -->\n",
       "<g id=\"node27\" class=\"node\">\n",
       "<title>5228709840</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"628.75\" cy=\"-483.5\" rx=\"79.5\" ry=\"18\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"628.75\" y=\"-478.45\" font-family=\"Times,serif\" font-size=\"14.00\">Sum{axes=None}</text>\n",
       "</g>\n",
       "<!-- 5228887168&#45;&gt;5228709840 -->\n",
       "<g id=\"edge28\" class=\"edge\">\n",
       "<title>5228887168-&gt;5228709840</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M567.17,-613.22C579.06,-607.83 591.3,-600.29 599.75,-590 617.64,-568.23 624.53,-536.26 627.17,-513.23\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"630.65,-513.57 628.08,-503.29 623.68,-512.93 630.65,-513.57\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"683.98\" y=\"-576.7\" font-family=\"Times,serif\" font-size=\"14.00\">Vector(float64, shape=(?,))</text>\n",
       "</g>\n",
       "<!-- 5228888032 -->\n",
       "<g id=\"node24\" class=\"node\">\n",
       "<title>5228888032</title>\n",
       "<ellipse fill=\"#ffaabb\" stroke=\"black\" cx=\"416.75\" cy=\"-429.5\" rx=\"62.61\" ry=\"18\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"416.75\" y=\"-424.45\" font-family=\"Times,serif\" font-size=\"14.00\">Cast{float64}</text>\n",
       "</g>\n",
       "<!-- 5224665776&#45;&gt;5228888032 -->\n",
       "<g id=\"edge25\" class=\"edge\">\n",
       "<title>5224665776-&gt;5228888032</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M408.83,-519.98C405.96,-514.32 403.2,-507.82 401.75,-501.5 398.16,-485.91 398.53,-481.17 401.75,-465.5 402.26,-463.03 402.95,-460.51 403.75,-458.03\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"406.94,-459.5 407.29,-448.91 400.41,-456.97 406.94,-459.5\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"470.75\" y=\"-478.45\" font-family=\"Times,serif\" font-size=\"14.00\">Vector(int64, shape=(1,))</text>\n",
       "</g>\n",
       "<!-- 5226719488 -->\n",
       "<g id=\"node25\" class=\"node\">\n",
       "<title>5226719488</title>\n",
       "<ellipse fill=\"#ffaaff\" stroke=\"black\" cx=\"416.75\" cy=\"-341\" rx=\"59.54\" ry=\"18\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"416.75\" y=\"-335.95\" font-family=\"Times,serif\" font-size=\"14.00\">Subtensor{i}</text>\n",
       "</g>\n",
       "<!-- 5228888032&#45;&gt;5226719488 -->\n",
       "<g id=\"edge26\" class=\"edge\">\n",
       "<title>5228888032-&gt;5226719488</title>\n",
       "<path fill=\"none\" stroke=\"blue\" d=\"M416.75,-411.41C416.75,-399.76 416.75,-384.05 416.75,-370.52\"/>\n",
       "<polygon fill=\"blue\" stroke=\"blue\" points=\"420.25,-370.86 416.75,-360.86 413.25,-370.86 420.25,-370.86\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"495.88\" y=\"-380.2\" font-family=\"Times,serif\" font-size=\"14.00\">0 Vector(float64, shape=(1,))</text>\n",
       "</g>\n",
       "<!-- 5228887456 -->\n",
       "<g id=\"node28\" class=\"node\">\n",
       "<title>5228887456</title>\n",
       "<ellipse fill=\"#ffaabb\" stroke=\"black\" cx=\"528.75\" cy=\"-252.5\" rx=\"44.7\" ry=\"18\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"528.75\" y=\"-247.45\" font-family=\"Times,serif\" font-size=\"14.00\">True_div</text>\n",
       "</g>\n",
       "<!-- 5226719488&#45;&gt;5228887456 -->\n",
       "<g id=\"edge30\" class=\"edge\">\n",
       "<title>5226719488-&gt;5228887456</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M427.35,-323.09C434.83,-312.18 445.65,-298.25 457.75,-288.5 466.41,-281.52 476.68,-275.46 486.69,-270.43\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"488.04,-273.67 495.57,-266.22 485.04,-267.34 488.04,-273.67\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"531.25\" y=\"-291.7\" font-family=\"Times,serif\" font-size=\"14.00\">1 Scalar(float64, shape=())</text>\n",
       "</g>\n",
       "<!-- 5228874832 -->\n",
       "<g id=\"node26\" class=\"node\">\n",
       "<title>5228874832</title>\n",
       "<polygon fill=\"green\" stroke=\"black\" points=\"336.25,-447.5 257.25,-447.5 257.25,-411.5 336.25,-411.5 336.25,-447.5\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"296.75\" y=\"-424.45\" font-family=\"Times,serif\" font-size=\"14.00\">val=0 int64</text>\n",
       "</g>\n",
       "<!-- 5228874832&#45;&gt;5226719488 -->\n",
       "<g id=\"edge27\" class=\"edge\">\n",
       "<title>5228874832-&gt;5226719488</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M320.75,-411.2C339.29,-397.84 365.1,-379.23 385.3,-364.67\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"387.12,-367.67 393.19,-358.98 383.03,-361.99 387.12,-367.67\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"370.87\" y=\"-380.2\" font-family=\"Times,serif\" font-size=\"14.00\">1</text>\n",
       "</g>\n",
       "<!-- 5228709840&#45;&gt;5228887456 -->\n",
       "<g id=\"edge29\" class=\"edge\">\n",
       "<title>5228709840-&gt;5228887456</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M631.76,-465.22C637.11,-428.7 644.3,-342.61 604.75,-288.5 597.8,-278.99 587.63,-271.99 576.97,-266.87\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"578.66,-263.79 568.09,-263.11 575.93,-270.23 578.66,-263.79\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"709.1\" y=\"-380.2\" font-family=\"Times,serif\" font-size=\"14.00\">0 Scalar(float64, shape=())</text>\n",
       "</g>\n",
       "<!-- 5224661792 -->\n",
       "<g id=\"node29\" class=\"node\">\n",
       "<title>5224661792</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"641,-182 416.5,-182 416.5,-146 641,-146 641,-182\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"528.75\" y=\"-158.95\" font-family=\"Times,serif\" font-size=\"14.00\">name=mean Scalar(float64, shape=())</text>\n",
       "</g>\n",
       "<!-- 5228887456&#45;&gt;5224661792 -->\n",
       "<g id=\"edge31\" class=\"edge\">\n",
       "<title>5228887456-&gt;5224661792</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M528.75,-234.41C528.75,-222.76 528.75,-207.05 528.75,-193.52\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"532.25,-193.86 528.75,-183.86 525.25,-193.86 532.25,-193.86\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"597\" y=\"-203.2\" font-family=\"Times,serif\" font-size=\"14.00\">Scalar(float64, shape=())</text>\n",
       "</g>\n",
       "<!-- 5228953680 -->\n",
       "<g id=\"node30\" class=\"node\">\n",
       "<title>5228953680</title>\n",
       "<ellipse fill=\"#ffaabb\" stroke=\"black\" cx=\"528.75\" cy=\"-91\" rx=\"27\" ry=\"18\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"528.75\" y=\"-85.95\" font-family=\"Times,serif\" font-size=\"14.00\">Neg</text>\n",
       "</g>\n",
       "<!-- 5224661792&#45;&gt;5228953680 -->\n",
       "<g id=\"edge32\" class=\"edge\">\n",
       "<title>5224661792-&gt;5228953680</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M528.75,-145.81C528.75,-138.23 528.75,-129.1 528.75,-120.54\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"532.25,-120.54 528.75,-110.54 525.25,-120.54 532.25,-120.54\"/>\n",
       "</g>\n",
       "<!-- 5228954352 -->\n",
       "<g id=\"node31\" class=\"node\">\n",
       "<title>5228954352</title>\n",
       "<polygon fill=\"blue\" stroke=\"black\" points=\"663.88,-36 393.62,-36 393.62,0 663.88,0 663.88,-36\"/>\n",
       "<text xml:space=\"preserve\" text-anchor=\"middle\" x=\"528.75\" y=\"-12.95\" font-family=\"Times,serif\" font-size=\"14.00\">name=cross_entropy Scalar(float64, shape=())</text>\n",
       "</g>\n",
       "<!-- 5228953680&#45;&gt;5228954352 -->\n",
       "<g id=\"edge33\" class=\"edge\">\n",
       "<title>5228953680-&gt;5228954352</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M528.75,-72.81C528.75,-65.23 528.75,-56.1 528.75,-47.54\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"532.25,-47.54 528.75,-37.54 525.25,-47.54 532.25,-47.54\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>"
      ],
      "text/plain": [
       "<IPython.core.display.SVG object>"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "SVG(pydotprint(loss, return_image=True, format=\"svg\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77b258df",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Comparsed to the text verison. It pays dividends to get used to reading these things!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "e793a2bc",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Neg [id A] 'cross_entropy'\n",
      " └─ True_div [id B] 'mean'\n",
      "    ├─ Sum{axes=None} [id C]\n",
      "    │  └─ Add [id D]\n",
      "    │     ├─ Mul [id E]\n",
      "    │     │  ├─ y [id F]\n",
      "    │     │  └─ Log [id G]\n",
      "    │     │     └─ Sigmoid [id H] 'p_class_0'\n",
      "    │     │        └─ Add [id I]\n",
      "    │     │           ├─ ExpandDims{axis=0} [id J]\n",
      "    │     │           │  └─ alpha [id K]\n",
      "    │     │           └─ Squeeze{axis=1} [id L]\n",
      "    │     │              └─ Blockwise{dot, (m,k),(k,n)->(m,n)} [id M]\n",
      "    │     │                 ├─ X [id N]\n",
      "    │     │                 └─ ExpandDims{axis=1} [id O]\n",
      "    │     │                    └─ beta [id P]\n",
      "    │     └─ Mul [id Q]\n",
      "    │        ├─ Sub [id R]\n",
      "    │        │  ├─ ExpandDims{axis=0} [id S]\n",
      "    │        │  │  └─ 1 [id T]\n",
      "    │        │  └─ y [id F]\n",
      "    │        └─ Log [id U]\n",
      "    │           └─ Sub [id V]\n",
      "    │              ├─ ExpandDims{axis=0} [id W]\n",
      "    │              │  └─ 1 [id X]\n",
      "    │              └─ Sigmoid [id H] 'p_class_0'\n",
      "    │                 └─ ···\n",
      "    └─ Subtensor{i} [id Y]\n",
      "       ├─ Cast{float64} [id Z]\n",
      "       │  └─ Shape [id BA]\n",
      "       │     └─ Add [id D]\n",
      "       │        └─ ···\n",
      "       └─ 0 [id BB]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x105c9e0e0>"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss.dprint()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33a6398f",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "To compute an update step, we will use a simple gradient descent algorithm, defined as:\n",
    "\n",
    "$$\\theta^\\prime = \\theta - \\eta \\nabla \\mathcal{L}(\\theta)$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "d24ad49c",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# The learning rate is another root variable we need to provide\n",
    "eta_pt = pt.tensor(\"learning_rate\", shape=())\n",
    "\n",
    "\n",
    "# We can compute gradients for a list of varibales too!\n",
    "# It's also good to apply some graph simplification before taking gradients. PyMC does this\n",
    "# by default, for example.\n",
    "d_alpha, d_beta = grad(\n",
    "    rewrite_graph(loss, include=(\"canonicalize\", \"stabilize\")), wrt=[alpha_pt, beta_pt]\n",
    ")\n",
    "\n",
    "# Apply gradient updates\n",
    "alpha_prime = alpha_pt - eta_pt * d_alpha\n",
    "beta_prime = beta_pt - eta_pt * d_beta"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e963af4",
   "metadata": {},
   "source": [
    "Compile our functions. We can make one for training, and one for prediction. Here we see another example of a PyTensor superpower: we don't have to know ahead of time how we want to use different outputs. We do all the steps of computation symbolically, then only at the end decide what will be used where. This type of thinking ahead is especially important when you're writing programs in Numba. You will need to think about what units of computation can be decomposed and jitted, for reuse in later jitted functions. We see that PyTensor does this reasoning for us.\n",
    "\n",
    "\n",
    "Getting to this point, we realized that we can output class probabilities, but not class labels. It's easy to make one more root variable for the prediction function, representing a prediction threshold for membership in class 1. Then we can make a `y_hat` variable and return it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "10903335",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "training_fn = pytensor.function(\n",
    "    [X_pt, y_pt, eta_pt, alpha_pt, beta_pt], [loss, alpha_prime, beta_prime]\n",
    ")\n",
    "\n",
    "threshold = pt.tensor(\"threshold\", shape=())\n",
    "y_hat = (p > threshold).astype(int)\n",
    "predict_fn = pytensor.function([X_pt, alpha_pt, beta_pt, threshold], [p, y_hat])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1480f470",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "It's worth having a look at the dprint of `training_fn` to see a couple features of pytensor we haven't run into yet.\n",
    "\n",
    "1. Sub-computations are automatically re-used! The first output, `Composite{...}.1 [id A] <Scalar(float64, shape=())> 9` is the cross entropy. It gets computed once (it's the first return, after all), but then it is also used in the 2nd return, to compute `alpha_prime`.\n",
    "2. `Composite{...}.0 [id I] <Vector(float64, shape=(?,))> 4` is the probability of class 0. That also gets computed once and re-used several times.\n",
    "3. Unused computations get truncated! Notice that `p_class_0` never appears as such. Instead, we got a composite sub-graph that compute that quantity given inputs. The thing itself was never requested, and is thus never used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "5e899f55",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Composite{...}.1 [id A] <Scalar(float64, shape=())> 'cross_entropy' d={0: [1]} 9\n",
      " ├─ Assert{msg='Could not broadcast dimensi...'} [id B] <Scalar(int64, shape=())> 8\n",
      " │  ├─ Shape_i{0} [id C] <Scalar(int64, shape=())> 0\n",
      " │  │  └─ X [id D] <Matrix(float64, shape=(?, ?))>\n",
      " │  └─ Eq [id E] <Scalar(bool, shape=())> 7\n",
      " │     ├─ Shape_i{0} [id C] <Scalar(int64, shape=())> 0\n",
      " │     │  └─ ···\n",
      " │     └─ Shape_i{0} [id F] <Scalar(int64, shape=())> 6\n",
      " │        └─ y [id G] <Vector(float64, shape=(?,))>\n",
      " └─ Sum{axes=None} [id H] <Scalar(float64, shape=())> 5\n",
      "    └─ Composite{...}.0 [id I] <Vector(float64, shape=(?,))> d={0: [1]} 4\n",
      "       ├─ ExpandDims{axis=0} [id J] <Vector(float64, shape=(1,))> 3\n",
      "       │  └─ alpha [id K] <Scalar(float64, shape=())>\n",
      "       ├─ CGemv{inplace} [id L] <Vector(float64, shape=(?,))> d={0: [0]} 2\n",
      "       │  ├─ AllocEmpty{dtype='float64'} [id M] <Vector(float64, shape=(?,))> 1\n",
      "       │  │  └─ Shape_i{0} [id C] <Scalar(int64, shape=())> 0\n",
      "       │  │     └─ ···\n",
      "       │  ├─ 1.0 [id N] <Scalar(float64, shape=())>\n",
      "       │  ├─ X [id D] <Matrix(float64, shape=(?, ?))>\n",
      "       │  ├─ beta [id O] <Vector(float64, shape=(?,))>\n",
      "       │  └─ 0.0 [id P] <Scalar(float64, shape=())>\n",
      "       └─ y [id G] <Vector(float64, shape=(?,))>\n",
      "Composite{(i2 - (i0 * i1))} [id Q] <Scalar(float64, shape=())> d={0: [1]} 13\n",
      " ├─ learning_rate [id R] <Scalar(float64, shape=())>\n",
      " ├─ Sum{axes=None} [id S] <Scalar(float64, shape=())> 12\n",
      " │  └─ Composite{((i2 / i1) + (i0 / i1))} [id T] <Vector(float64, shape=(?,))> d={0: [0]} 11\n",
      " │     ├─ Composite{...}.1 [id I] <Vector(float64, shape=(?,))> d={0: [1]} 4\n",
      " │     │  └─ ···\n",
      " │     ├─ ExpandDims{axis=0} [id U] <Vector(float64, shape=(1,))> 10\n",
      " │     │  └─ Composite{...}.0 [id A] <Scalar(float64, shape=())> d={0: [1]} 9\n",
      " │     │     └─ ···\n",
      " │     └─ Composite{...}.2 [id I] <Vector(float64, shape=(?,))> d={0: [1]} 4\n",
      " │        └─ ···\n",
      " └─ alpha [id K] <Scalar(float64, shape=())>\n",
      "CGemv{no_inplace} [id V] <Vector(float64, shape=(?,))> 16\n",
      " ├─ beta [id O] <Vector(float64, shape=(?,))>\n",
      " ├─ Neg [id W] <Scalar(float64, shape=())> 15\n",
      " │  └─ learning_rate [id R] <Scalar(float64, shape=())>\n",
      " ├─ Transpose{axes=[1, 0]} [id X] <Matrix(float64, shape=(?, ?))> 'X.T' 14\n",
      " │  └─ X [id D] <Matrix(float64, shape=(?, ?))>\n",
      " ├─ Composite{((i2 / i1) + (i0 / i1))} [id T] <Vector(float64, shape=(?,))> d={0: [0]} 11\n",
      " │  └─ ···\n",
      " └─ 1.0 [id N] <Scalar(float64, shape=())>\n",
      "\n",
      "Inner graphs:\n",
      "\n",
      "Composite{...} [id A] d={0: [1]}\n",
      " ← Cast{float64} [id Y] <float64> 'o0'\n",
      "    └─ i0 [id Z] <int64>\n",
      " ← neg [id BA] <float64> 'o1'\n",
      "    └─ true_div [id BB] <float64>\n",
      "       ├─ i1 [id BC] <float64>\n",
      "       └─ Cast{float64} [id Y] <float64> 'o0'\n",
      "          └─ ···\n",
      "\n",
      "Composite{...} [id I] d={0: [1]}\n",
      " ← add [id BD] <float64> 'o0'\n",
      "    ├─ mul [id BE] <float64>\n",
      "    │  ├─ -1.0 [id BF] <float64>\n",
      "    │  ├─ i2 [id BG] <float64>\n",
      "    │  └─ scalar_softplus [id BH] <float64>\n",
      "    │     └─ neg [id BI] <float64> 't4'\n",
      "    │        └─ add [id BJ] <float64> 't14'\n",
      "    │           ├─ i0 [id BK] <float64>\n",
      "    │           └─ i1 [id BL] <float64>\n",
      "    └─ mul [id BM] <float64>\n",
      "       ├─ -1.0 [id BF] <float64>\n",
      "       ├─ sub [id BN] <float64> 't7'\n",
      "       │  ├─ 1.0 [id BO] <float64>\n",
      "       │  └─ i2 [id BG] <float64>\n",
      "       └─ scalar_softplus [id BP] <float64>\n",
      "          └─ add [id BJ] <float64> 't14'\n",
      "             └─ ···\n",
      " ← mul [id BQ] <float64> 'o1'\n",
      "    ├─ sub [id BN] <float64> 't7'\n",
      "    │  └─ ···\n",
      "    └─ sigmoid [id BR] <float64>\n",
      "       └─ add [id BJ] <float64> 't14'\n",
      "          └─ ···\n",
      " ← mul [id BS] <float64> 'o2'\n",
      "    ├─ -1.0 [id BF] <float64>\n",
      "    ├─ i2 [id BG] <float64>\n",
      "    └─ sigmoid [id BT] <float64>\n",
      "       └─ neg [id BI] <float64> 't4'\n",
      "          └─ ···\n",
      "\n",
      "Composite{(i2 - (i0 * i1))} [id Q] d={0: [1]}\n",
      " ← sub [id BU] <float64> 'o0'\n",
      "    ├─ i2 [id BV] <float64>\n",
      "    └─ mul [id BW] <float64>\n",
      "       ├─ i0 [id BX] <float64>\n",
      "       └─ i1 [id BY] <float64>\n",
      "\n",
      "Composite{((i2 / i1) + (i0 / i1))} [id T] d={0: [0]}\n",
      " ← add [id BZ] <float64> 'o0'\n",
      "    ├─ true_div [id CA] <float64>\n",
      "    │  ├─ i2 [id CB] <float64>\n",
      "    │  └─ i1 [id CC] <float64>\n",
      "    └─ true_div [id CD] <float64>\n",
      "       ├─ i0 [id CE] <float64>\n",
      "       └─ i1 [id CC] <float64>\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x105c9e0e0>"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "training_fn.dprint(print_type=True, print_destroy_map=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c880ffa",
   "metadata": {},
   "source": [
    "The last point about truncation is even more clear when we look at the graph for `predict_fn`. Now the loss is nowhere to be seen! We just compute exactly what was requested.\n",
    "\n",
    "Note that we also have a nice example of re-using computation. We compute `p_class_0` exactly once, then re-use it in the 2nd output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "1a50f419",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Composite{...}.0 [id A] <Vector(float64, shape=(?,))> 'p_class_0' 5\n",
      " ├─ ExpandDims{axis=0} [id B] <Vector(float64, shape=(1,))> 4\n",
      " │  └─ alpha [id C] <Scalar(float64, shape=())>\n",
      " ├─ CGemv{inplace} [id D] <Vector(float64, shape=(?,))> 3\n",
      " │  ├─ AllocEmpty{dtype='float64'} [id E] <Vector(float64, shape=(?,))> 2\n",
      " │  │  └─ Shape_i{0} [id F] <Scalar(int64, shape=())> 1\n",
      " │  │     └─ X [id G] <Matrix(float64, shape=(?, ?))>\n",
      " │  ├─ 1.0 [id H] <Scalar(float64, shape=())>\n",
      " │  ├─ X [id G] <Matrix(float64, shape=(?, ?))>\n",
      " │  ├─ beta [id I] <Vector(float64, shape=(?,))>\n",
      " │  └─ 0.0 [id J] <Scalar(float64, shape=())>\n",
      " └─ ExpandDims{axis=0} [id K] <Vector(float64, shape=(1,))> 0\n",
      "    └─ threshold [id L] <Scalar(float64, shape=())>\n",
      "Composite{...}.1 [id A] <Vector(int64, shape=(?,))> 5\n",
      " └─ ···\n",
      "\n",
      "Inner graphs:\n",
      "\n",
      "Composite{...} [id A]\n",
      " ← sigmoid [id M] <float64> 'o0'\n",
      "    └─ add [id N] <float64>\n",
      "       ├─ i0 [id O] <float64>\n",
      "       └─ i1 [id P] <float64>\n",
      " ← Cast{int64} [id Q] <int64> 'o1'\n",
      "    └─ GT [id R] <bool>\n",
      "       ├─ sigmoid [id M] <float64> 'o0'\n",
      "       │  └─ ···\n",
      "       └─ i2 [id S] <float64>\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x105c9e0e0>"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predict_fn.dprint(print_type=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "615f13ef",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "b57e8e42",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from sklearn.datasets import make_classification\n",
    "\n",
    "\n",
    "seed = sum(map(ord, \"I <3 Pytensor\"))\n",
    "rng = np.random.default_rng(seed)\n",
    "\n",
    "X, y = make_classification(\n",
    "    n_samples=100,\n",
    "    n_features=5,\n",
    "    n_redundant=0,\n",
    "    n_informative=5,\n",
    "    n_classes=2,\n",
    "    random_state=seed,\n",
    ")\n",
    "\n",
    "# Initial values for parameteres\n",
    "beta = rng.normal(0, 1, size=(5,))\n",
    "alpha = rng.normal(0, 1)\n",
    "\n",
    "learning_rate = 1e-1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "f4400eb1",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "n_steps = 1_000\n",
    "histories = [np.empty(n_steps), np.empty(n_steps), np.empty((n_steps, 5))]\n",
    "\n",
    "for t in range(n_steps):\n",
    "    loss_val, alpha, beta = training_fn(X, y, learning_rate, alpha, beta)\n",
    "    histories[0][t] = loss_val\n",
    "    histories[1][t] = alpha\n",
    "    histories[2][t] = beta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "58052478",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 149660 characters]",
      "text/plain": [
       "<Figure size 800x1000 with 3 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 1011,
       "width": 811
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(\n",
    "    nrows=3,\n",
    "    ncols=1,\n",
    "    figsize=(8, 10),\n",
    "    sharex=True,\n",
    "    layout=\"constrained\",\n",
    ")\n",
    "for axis, data, name in zip(\n",
    "    fig.axes,\n",
    "    histories,\n",
    "    [\"cross-entropy\", \"alpha\", \"beta\"],\n",
    "):\n",
    "    axis.plot(data)\n",
    "    axis.set(title=name)\n",
    "\n",
    "axis.set(xlabel=\"steps\")\n",
    "\n",
    "fig.suptitle(\"Training\", fontsize=18, fontweight=\"bold\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af29f666",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "`make_classification` doesn't return true parameters to check against, so we'll just look at the confusion matrix. Seems like we learned something. It's a toy problem anyway!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "cacdee5b",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[36, 14],\n",
       "       [11, 39]])"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from sklearn.metrics import confusion_matrix\n",
    "\n",
    "\n",
    "confusion_matrix(y, predict_fn(X, alpha, beta, 0.5)[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7fbb667d",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Taking a step back\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1338584",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### How does it compare with alternative frameworks\n",
    "\n",
    "* Graph is built explicitly with placeholder inputs (common source of confusion for users)\n",
    "* It is focused on array (tensor) operations (dense and sparse). Tries to look almost like numpy / scipy, (until the abstraction breaks).\n",
    "  * There is narrow / hidden support for other types like scalars, lists, slices, random Generators, strings, None (although easy to extend)\n",
    "* Functional design (there is no variable mutation when defining graphs)\n",
    "* Strong focus on hackability / graph manipulation\n",
    "* Evolved from:\n",
    "  1. Theano which strongly inspired Tensorflow 1.x and JAX. Many concepts stood the test of time. Others have aged and provide some drag.\n",
    "  2. Aesara, which cleaned up the codebase, added alternative backends (Numba and JAX) and proved there's some interest out there in a library like this."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12420067",
   "metadata": {},
   "source": [
    "We hope you enjoyed this introduction to PyTensor. This is just the beginning of what you can do with it. Please explore the documentation and the gallery for more examples and applications."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "556202cf",
   "metadata": {},
   "source": [
    "## Authors\n",
    "\n",
    "- Jesse Grabowski and Ricardo Vieira in September 2025\n",
    "- Ported to the repository by Juan Orduz in September 2025"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f06cf7a2",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    ":::{bibliography} :filter: docname in docnames"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b3d194f",
   "metadata": {},
   "source": [
    "## Watermark "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "8277fb12",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Last updated: Sat Sep 13 2025\n",
      "\n",
      "Python implementation: CPython\n",
      "Python version       : 3.12.11\n",
      "IPython version      : 9.5.0\n",
      "\n",
      "pytensor: 2.31.7\n",
      "\n",
      "pytensor  : 2.31.7\n",
      "graphviz  : 0.21\n",
      "pymc      : 5.25.1\n",
      "sklearn   : 1.7.2\n",
      "scipy     : 1.16.2\n",
      "matplotlib: 3.10.6\n",
      "numpy     : 2.2.6\n",
      "IPython   : 9.5.0\n",
      "\n",
      "Watermark: 2.5.0\n",
      "\n"
     ]
    }
   ],
   "source": [
    "%load_ext watermark\n",
    "%watermark -n -u -v -iv -w -p pytensor"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Slideshow",
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
