notebooks/walkthrough/graph_transformation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/gist/ricardoV94/65a559e6b977c2e3709ad636e22f083d/graph_transformation.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "!pip install 'pytensor>=2.28.3'\n",
    "try:\n",
    "    import pytensor_workshop\n",
    "except ModuleNotFoundError:\n",
    "    !pip install git+https://github.com/pymc-devs/pytensor-workshop.git"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yOqTXAc5mz3V"
   },
   "source": [
    "## Graph transformations: two perspectives.\n",
    "\n",
    "One of the most powerful uses of pytensor is the transformation of one graph to another. In this Notebook we will explore two transformations with a different flavor: gradient and graph subset (and shape as a bonus)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "r3SssvMwnHWa"
   },
   "source": [
    "## Gradient transformation\n",
    "\n",
    "`pytensor.gradient.grad` eagerly generates a graph that represents the gradient of a scalar function with respect to some inputs specified by the user. Let's see it in action, and then recreate it ourselves"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-11T10:51:53.841948Z",
     "start_time": "2025-03-11T10:51:52.980751Z"
    },
    "id": "mq4GFMA0evQn"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import pytensor\n",
    "import pytensor.tensor as pt\n",
    "from pytensor.graph import rewrite_graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-11T10:51:54.035319Z",
     "start_time": "2025-03-11T10:51:54.018446Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "F_O4ua59e4Rc",
    "outputId": "b591e683-0955-4f37-effe-03c3cb6ef6b0"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sum{axes=None} [id A] shape=()\n",
      " └─ Add [id B] shape=(3, 2)\n",
      "    ├─ Exp [id C] shape=(3, 2)\n",
      "    │  └─ Transpose{axes=[1, 0]} [id D] shape=(3, 2) 'x.T'\n",
      "    │     └─ x [id E] shape=(2, 3)\n",
      "    └─ Sqr [id F] shape=(3, 2)\n",
      "       └─ Transpose{axes=[1, 0]} [id D] shape=(3, 2) 'x.T'\n",
      "          └─ ···\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x7fe53b434d60>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = pt.matrix(\"x\", shape=(2, 3))\n",
    "xT = x.T\n",
    "y = (pt.exp(xT) + pt.square(xT)).sum()\n",
    "y.dprint(print_shape=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-11T10:51:54.225240Z",
     "start_time": "2025-03-11T10:51:54.202678Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "A0MxfRP_gYqc",
    "outputId": "9b697c85-d601-4553-ccf3-adf09df54499"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Transpose{axes=[1, 0]} [id A] shape=(2, 3)\n",
      " └─ Add [id B] shape=(3, 2)\n",
      "    ├─ Mul [id C] shape=(3, 2)\n",
      "    │  ├─ Second [id D] shape=(3, 2)\n",
      "    │  │  ├─ Add [id E] shape=(3, 2)\n",
      "    │  │  │  ├─ Exp [id F] shape=(3, 2)\n",
      "    │  │  │  │  └─ Transpose{axes=[1, 0]} [id G] shape=(3, 2) 'x.T'\n",
      "    │  │  │  │     └─ x [id H] shape=(2, 3)\n",
      "    │  │  │  └─ Sqr [id I] shape=(3, 2)\n",
      "    │  │  │     └─ Transpose{axes=[1, 0]} [id G] shape=(3, 2) 'x.T'\n",
      "    │  │  │        └─ ···\n",
      "    │  │  └─ ExpandDims{axes=[0, 1]} [id J] shape=(1, 1)\n",
      "    │  │     └─ Second [id K] shape=()\n",
      "    │  │        ├─ Sum{axes=None} [id L] shape=()\n",
      "    │  │        │  └─ Add [id E] shape=(3, 2)\n",
      "    │  │        │     └─ ···\n",
      "    │  │        └─ 1.0 [id M] shape=()\n",
      "    │  └─ Exp [id N] shape=(3, 2)\n",
      "    │     └─ Transpose{axes=[1, 0]} [id G] shape=(3, 2) 'x.T'\n",
      "    │        └─ ···\n",
      "    └─ Mul [id O] shape=(3, 2)\n",
      "       ├─ Mul [id P] shape=(3, 2)\n",
      "       │  ├─ Second [id D] shape=(3, 2)\n",
      "       │  │  └─ ···\n",
      "       │  └─ Transpose{axes=[1, 0]} [id G] shape=(3, 2) 'x.T'\n",
      "       │     └─ ···\n",
      "       └─ ExpandDims{axes=[0, 1]} [id Q] shape=(1, 1)\n",
      "          └─ 2 [id R] shape=()\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x7fe53b434d60>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pt.grad(y, wrt=x).dprint(print_shape=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-11T10:51:55.661139Z",
     "start_time": "2025-03-11T10:51:54.410381Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "NR1ZqSygfH5r",
    "outputId": "b5f5b86c-1fc8-4cd3-e7ad-55216fa34ab0"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Add [id A]\n",
      " ├─ Exp [id B]\n",
      " │  └─ x [id C]\n",
      " └─ Mul [id D]\n",
      "    ├─ [[2.]] [id E]\n",
      "    └─ x [id C]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x7fe53b434d60>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rewrite_graph(pt.grad(y, wrt=x)).dprint()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-11T10:51:55.699260Z",
     "start_time": "2025-03-11T10:51:55.690593Z"
    },
    "id": "6_xpfvEkfkt1"
   },
   "outputs": [],
   "source": [
    "from pytensor.graph.basic import io_toposort\n",
    "\n",
    "def didactic_grad(output, wrt):\n",
    "    if output.type.ndim != 0:\n",
    "        raise ValueError(\"gradient output must be scalar\")\n",
    "\n",
    "    if not isinstance(wrt, tuple | list):\n",
    "        wrt = [wrt]\n",
    "\n",
    "    # The vectors in vector-jacobian-products\n",
    "    acc_cotangents = {output: pt.as_tensor(1.0)}\n",
    "\n",
    "    # Go in reverse topological order from output to inputs\n",
    "    for i, node in enumerate(reversed(io_toposort(wrt, [output]))):\n",
    "        print(i)\n",
    "        node.dprint(depth=2, print_shape=True)\n",
    "        print()\n",
    "        cotangents = [acc_cotangents.get(output, None) for output in node.outputs]\n",
    "\n",
    "        input_cotangents = node.op.L_op(node.inputs, node.outputs, cotangents)\n",
    "\n",
    "        for input, input_cotangent in zip(node.inputs, input_cotangents, strict=True):\n",
    "            if input_cotangent is None:\n",
    "                # Input is disconnected from the gradient\n",
    "                continue\n",
    "\n",
    "            if input not in acc_cotangents:\n",
    "                acc_cotangents[input] = input_cotangent\n",
    "            else:\n",
    "                acc_cotangents[input] += input_cotangent\n",
    "\n",
    "    return [acc_cotangents[var] for var in wrt]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-11T10:51:55.841674Z",
     "start_time": "2025-03-11T10:51:55.817738Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "bZfl6VKlpfiL",
    "outputId": "cceb06ec-5fe6-4bf0-ad52-72de2ada7734"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "Sum{axes=None} [id A] shape=()\n",
      " └─ Add [id B] shape=(3, 2)\n",
      "\n",
      "1\n",
      "Add [id A] shape=(3, 2)\n",
      " ├─ Exp [id B] shape=(3, 2)\n",
      " └─ Sqr [id C] shape=(3, 2)\n",
      "\n",
      "2\n",
      "Exp [id A] shape=(3, 2)\n",
      " └─ Transpose{axes=[1, 0]} [id B] shape=(3, 2) 'x.T'\n",
      "\n",
      "3\n",
      "Sqr [id A] shape=(3, 2)\n",
      " └─ Transpose{axes=[1, 0]} [id B] shape=(3, 2) 'x.T'\n",
      "\n",
      "4\n",
      "Transpose{axes=[1, 0]} [id A] shape=(3, 2) 'x.T'\n",
      " └─ x [id B] shape=(2, 3)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "[grad_y] = didactic_grad(y, wrt=x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-11T10:51:55.956845Z",
     "start_time": "2025-03-11T10:51:55.948582Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "5OmLx5CWqyIq",
    "outputId": "8d1c6558-d196-4934-c85c-d1c0eb13baa3"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Transpose{axes=[1, 0]} [id A]\n",
      " └─ Add [id B]\n",
      "    ├─ Mul [id C]\n",
      "    │  ├─ Second [id D]\n",
      "    │  │  ├─ Add [id E]\n",
      "    │  │  │  ├─ Exp [id F]\n",
      "    │  │  │  │  └─ Transpose{axes=[1, 0]} [id G] 'x.T'\n",
      "    │  │  │  │     └─ x [id H]\n",
      "    │  │  │  └─ Sqr [id I]\n",
      "    │  │  │     └─ Transpose{axes=[1, 0]} [id G] 'x.T'\n",
      "    │  │  │        └─ ···\n",
      "    │  │  └─ ExpandDims{axes=[0, 1]} [id J]\n",
      "    │  │     └─ 1.0 [id K]\n",
      "    │  └─ Exp [id L]\n",
      "    │     └─ Transpose{axes=[1, 0]} [id G] 'x.T'\n",
      "    │        └─ ···\n",
      "    └─ Mul [id M]\n",
      "       ├─ Mul [id N]\n",
      "       │  ├─ Second [id D]\n",
      "       │  │  └─ ···\n",
      "       │  └─ Transpose{axes=[1, 0]} [id G] 'x.T'\n",
      "       │     └─ ···\n",
      "       └─ ExpandDims{axes=[0, 1]} [id O]\n",
      "          └─ 2 [id P]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x7fe53b434d60>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "grad_y.dprint()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-11T10:51:56.041824Z",
     "start_time": "2025-03-11T10:51:56.011450Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "JibDgMx9sOeq",
    "outputId": "9a6e1c2a-3d64-497f-a15f-0bf7e3a0ba52"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Add [id A]\n",
      " ├─ Exp [id B]\n",
      " │  └─ x [id C]\n",
      " └─ Mul [id D]\n",
      "    ├─ [[2.]] [id E]\n",
      "    └─ x [id C]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x7fe53b434d60>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rewrite_graph(grad_y).dprint()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2MCaf3SJuvXD"
   },
   "source": [
    "The graph transformations can be implemented in a simple manner, because we always have the rewrite machinery to cleanup things after. This is a recurring motive in PyTensor!\n",
    "\n",
    "You can argue that it's messier, but it also allows code to be modular and readable. Each Op only has to know how to differentiate itself in a manner that's correct. Efficiency is left for whole graph rewriting"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gwl7HN-Fves2"
   },
   "source": [
    "### Aside: Let's see the graph rewriting step by step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-11T10:51:56.125689Z",
     "start_time": "2025-03-11T10:51:56.094151Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "3tHAzymmuuD5",
    "outputId": "a5af30fa-b5d6-4cb7-c68b-02218f8dc57f"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rewriting: rewrite local_dimshuffle_lift replaces Transpose{axes=[1, 0]}.0 of Transpose{axes=[1, 0]}(Add.0) with Add.0 of Add(Mul.0, Mul.0)\n",
      "rewriting: rewrite local_mul_canonizer replaces Mul.0 of Mul(ExpandDims{axes=[0, 1]}.0, Exp.0) with Exp.0 of Exp(x)\n",
      "rewriting: rewrite local_mul_canonizer replaces Mul.0 of Mul(Mul.0, ExpandDims{axes=[0, 1]}.0) with Mul.0 of Mul(ExpandDims{axes=[0, 1]}.0, x)\n"
     ]
    }
   ],
   "source": [
    "with pytensor.config.change_flags(optimizer_verbose=True):\n",
    "    rewrite_graph(grad_y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "957AHYKwvtGw"
   },
   "source": [
    "Sounds reasonable but what is actually going on?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-11T10:51:56.317813Z",
     "start_time": "2025-03-11T10:51:56.270110Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "N3VBnk6cwIqr",
    "outputId": "9fff01a8-e66b-4aec-f35f-583b6bdab00f"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Transpose{axes=[1, 0]} [id A] shape=(2, 3) 8\n",
      " └─ Add [id B] shape=(3, 2) 7\n",
      "    ├─ Mul [id C] shape=(3, 2) 6\n",
      "    │  ├─ ExpandDims{axes=[0, 1]} [id D] shape=(1, 1) 2\n",
      "    │  │  └─ 1.0 [id E] shape=()\n",
      "    │  └─ Exp [id F] shape=(3, 2) 5\n",
      "    │     └─ Transpose{axes=[1, 0]} [id G] shape=(3, 2) 'x.T' 1\n",
      "    │        └─ x [id H] shape=(2, 3)\n",
      "    └─ Mul [id I] shape=(3, 2) 4\n",
      "       ├─ Mul [id J] shape=(3, 2) 3\n",
      "       │  ├─ ExpandDims{axes=[0, 1]} [id D] shape=(1, 1) 2\n",
      "       │  │  └─ ···\n",
      "       │  └─ Transpose{axes=[1, 0]} [id G] shape=(3, 2) 'x.T' 1\n",
      "       │     └─ ···\n",
      "       └─ ExpandDims{axes=[0, 1]} [id K] shape=(1, 1) 0\n",
      "          └─ 2 [id L] shape=()\n",
      "\n",
      ">>> local_dimshuffle_lift\n",
      "Add [id A] shape=(2, 3) 7\n",
      " ├─ Mul [id B] shape=(2, 3) 6\n",
      " │  ├─ ExpandDims{axes=[0, 1]} [id C] shape=(1, 1) 5\n",
      " │  │  └─ 1.0 [id D] shape=()\n",
      " │  └─ Exp [id E] shape=(2, 3) 4\n",
      " │     └─ x [id F] shape=(2, 3)\n",
      " └─ Mul [id G] shape=(2, 3) 3\n",
      "    ├─ Mul [id H] shape=(2, 3) 2\n",
      "    │  ├─ ExpandDims{axes=[0, 1]} [id I] shape=(1, 1) 1\n",
      "    │  │  └─ 1.0 [id D] shape=()\n",
      "    │  └─ x [id F] shape=(2, 3)\n",
      "    └─ ExpandDims{axes=[0, 1]} [id J] shape=(1, 1) 0\n",
      "       └─ 2 [id K] shape=()\n",
      "\n",
      ">>> local_mul_canonizer\n",
      "Add [id A] shape=(2, 3) 5\n",
      " ├─ Exp [id B] shape=(2, 3) 4\n",
      " │  └─ x [id C] shape=(2, 3)\n",
      " └─ Mul [id D] shape=(2, 3) 3\n",
      "    ├─ Mul [id E] shape=(2, 3) 2\n",
      "    │  ├─ ExpandDims{axes=[0, 1]} [id F] shape=(1, 1) 1\n",
      "    │  │  └─ 1.0 [id G] shape=()\n",
      "    │  └─ x [id C] shape=(2, 3)\n",
      "    └─ ExpandDims{axes=[0, 1]} [id H] shape=(1, 1) 0\n",
      "       └─ 2 [id I] shape=()\n",
      "\n",
      ">>> local_mul_canonizer\n",
      "Add [id A] shape=(2, 3) 3\n",
      " ├─ Exp [id B] shape=(2, 3) 2\n",
      " │  └─ x [id C] shape=(2, 3)\n",
      " └─ Mul [id D] shape=(2, 3) 1\n",
      "    ├─ ExpandDims{axes=[0, 1]} [id E] shape=(1, 1) 0\n",
      "    │  └─ 2.0 [id F] shape=()\n",
      "    └─ x [id C] shape=(2, 3)\n",
      "\n",
      ">>> constant_folding\n",
      "Add [id A] shape=(2, 3) 2\n",
      " ├─ Exp [id B] shape=(2, 3) 1\n",
      " │  └─ x [id C] shape=(2, 3)\n",
      " └─ Mul [id D] shape=(2, 3) 0\n",
      "    ├─ [[2.]] [id E] shape=(1, 1)\n",
      "    └─ x [id C] shape=(2, 3)\n"
     ]
    }
   ],
   "source": [
    "from pytensor.graph import FunctionGraph\n",
    "from pytensor_workshop import FullHistory\n",
    "\n",
    "fg = FunctionGraph(outputs=[grad_y])\n",
    "history = FullHistory()\n",
    "fg.attach_feature(history)\n",
    "\n",
    "rewrite_graph(fg)\n",
    "\n",
    "# Replay rewrites\n",
    "history.start()\n",
    "pytensor.dprint(fg, print_shape=True)\n",
    "with pytensor.config.change_flags(optimizer_verbose = True):\n",
    "    for i in range(4):\n",
    "        print()\n",
    "        print(\">>> \", end=\"\")\n",
    "        pytensor.dprint(history.next(), print_shape=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aBp_DN4G3FKy"
   },
   "source": [
    "## Graph subset\n",
    "\n",
    "Sometimes we are only interested in a subset of a variable. Perhaps we have a training function that computes the outcome over a large number of days, but for prediction we only to compute the outcome for a single day.\n",
    "\n",
    "We could have a graph transformation like `grad`, but for these purposes we'll take a more indirect route. We'll specify what we want and let PyTensor provide the best solution to that problem via rewrites."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-11T10:51:56.399058Z",
     "start_time": "2025-03-11T10:51:56.383460Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "HXVOWljq0rUk",
    "outputId": "3edaf9ab-8d4c-4087-8fb5-daf2d0c03fc2"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Blockwise{dot, (m,k),(k,n)->(m,n)} [id A] shape=(512, 256)\n",
      " ├─ Cos [id B] shape=(512, 64)\n",
      " │  └─ x [id C] shape=(512, 64)\n",
      " └─ Exp [id D] shape=(64, 256)\n",
      "    └─ y [id E] shape=(64, 256)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x7fe53b434d60>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = pt.matrix(\"x\", shape=(512, 64))\n",
    "y = pt.matrix(\"y\", shape=(64, 256))\n",
    "\n",
    "outs = (pt.cos(x) @ pt.exp(y))\n",
    "outs.dprint(print_shape=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "W48MCJ0_3v2q"
   },
   "source": [
    "Let's say we only need the last entry of the first row of this function. Perhaps the function was defined by the user so it's not possible for us to just redefine it. Or we want to avoid mistakes in our rewrite\n",
    "\n",
    "Let's just tell PyTensor what we need."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-11T10:51:56.474495Z",
     "start_time": "2025-03-11T10:51:56.466455Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "QUXipgy03t3o",
    "outputId": "8ce419a1-c603-4410-c516-d40aa966ec53"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Subtensor{i, j} [id A] shape=()\n",
      " ├─ Blockwise{dot, (m,k),(k,n)->(m,n)} [id B] shape=(512, 256)\n",
      " │  ├─ Cos [id C] shape=(512, 64)\n",
      " │  │  └─ x [id D] shape=(512, 64)\n",
      " │  └─ Exp [id E] shape=(64, 256)\n",
      " │     └─ y [id F] shape=(64, 256)\n",
      " ├─ 0 [id G] shape=()\n",
      " └─ -1 [id H] shape=()\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x7fe53b434d60>"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out = outs[0, -1]\n",
    "out.dprint(print_shape=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QQawxhlo36iQ"
   },
   "source": [
    "We added an indexing operation. This is not a graph transformation, at least not yet! It's just an operation that takes a pre-computed variable and selects specific entries.\n",
    "\n",
    "If we were to evaluate this function without any further rewrites we would still compute all 512 rows and 256 columns, to then discand everything but the entry we need.\n",
    "\n",
    "Can PyTensor figure out something better?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-11T10:51:56.565393Z",
     "start_time": "2025-03-11T10:51:56.547456Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "d5tKosTn2mqU",
    "outputId": "2fd2672b-097a-4a15-a534-f71002df6de4"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dot [id A] shape=()\n",
      " ├─ Cos [id B] shape=(64,)\n",
      " │  └─ Subtensor{i} [id C] shape=(64,)\n",
      " │     ├─ x [id D] shape=(512, 64)\n",
      " │     └─ 0 [id E] shape=()\n",
      " └─ Exp [id F] shape=(64,)\n",
      "    └─ Subtensor{:, i} [id G] shape=(64,)\n",
      "       ├─ y [id H] shape=(64, 256)\n",
      "       └─ -1 [id I] shape=()\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x7fe53b434d60>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rewrite_graph(out).dprint(print_shape=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6mIvw2JE4cOH"
   },
   "source": [
    "Much better! Not only we only compute one vector product, we also only compute the cosine and exp for the relevant rows and columns"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Hp9YxArc4v29"
   },
   "source": [
    "### How did it get there?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-11T10:51:56.646601Z",
     "start_time": "2025-03-11T10:51:56.634993Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "2aDvQjvj4YDD",
    "outputId": "2ab1253a-e70d-44b4-bb3b-5dea275c1f37"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rewriting: rewrite local_subtensor_of_dot replaces Subtensor{i, j}.0 of Subtensor{i, j}(dot.0, 0, -1) with dot.0 of dot(Subtensor{i}.0, Subtensor{:, i}.0)\n",
      "rewriting: rewrite local_subtensor_lift replaces Subtensor{i}.0 of Subtensor{i}(Cos.0, 0) with Cos.0 of Cos(Subtensor{i}.0)\n",
      "rewriting: rewrite local_subtensor_lift replaces Subtensor{:, i}.0 of Subtensor{:, i}(Exp.0, -1) with Exp.0 of Exp(Subtensor{:, i}.0)\n"
     ]
    }
   ],
   "source": [
    "with pytensor.config.change_flags(optimizer_verbose=True):\n",
    "    rewrite_graph(out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-11T10:51:56.733387Z",
     "start_time": "2025-03-11T10:51:56.700729Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "7RPTeH-12yTB",
    "outputId": "1865a6f4-c134-4869-c6bf-3ab19342bf61"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Subtensor{i, j} [id A] shape=() 3\n",
      " ├─ dot [id B] shape=(512, 256) 2\n",
      " │  ├─ Cos [id C] shape=(512, 64) 1\n",
      " │  │  └─ x [id D] shape=(512, 64)\n",
      " │  └─ Exp [id E] shape=(64, 256) 0\n",
      " │     └─ y [id F] shape=(64, 256)\n",
      " ├─ 0 [id G] shape=()\n",
      " └─ -1 [id H] shape=()\n",
      "\n",
      ">>> local_subtensor_of_dot\n",
      "dot [id A] shape=() 4\n",
      " ├─ Subtensor{i} [id B] shape=(64,) 3\n",
      " │  ├─ Cos [id C] shape=(512, 64) 2\n",
      " │  │  └─ x [id D] shape=(512, 64)\n",
      " │  └─ 0 [id E] shape=()\n",
      " └─ Subtensor{:, i} [id F] shape=(64,) 1\n",
      "    ├─ Exp [id G] shape=(64, 256) 0\n",
      "    │  └─ y [id H] shape=(64, 256)\n",
      "    └─ -1 [id I] shape=()\n",
      "\n",
      ">>> local_subtensor_lift\n",
      "dot [id A] shape=() 4\n",
      " ├─ Cos [id B] shape=(64,) 3\n",
      " │  └─ Subtensor{i} [id C] shape=(64,) 2\n",
      " │     ├─ x [id D] shape=(512, 64)\n",
      " │     └─ 0 [id E] shape=()\n",
      " └─ Subtensor{:, i} [id F] shape=(64,) 1\n",
      "    ├─ Exp [id G] shape=(64, 256) 0\n",
      "    │  └─ y [id H] shape=(64, 256)\n",
      "    └─ -1 [id I] shape=()\n",
      "\n",
      ">>> local_subtensor_lift\n",
      "dot [id A] shape=() 4\n",
      " ├─ Cos [id B] shape=(64,) 3\n",
      " │  └─ Subtensor{i} [id C] shape=(64,) 2\n",
      " │     ├─ x [id D] shape=(512, 64)\n",
      " │     └─ 0 [id E] shape=()\n",
      " └─ Exp [id F] shape=(64,) 1\n",
      "    └─ Subtensor{:, i} [id G] shape=(64,) 0\n",
      "       ├─ y [id H] shape=(64, 256)\n",
      "       └─ -1 [id I] shape=()\n"
     ]
    }
   ],
   "source": [
    "# Now step by step\n",
    "from pytensor.graph import FunctionGraph\n",
    "\n",
    "fg = FunctionGraph(outputs=[out])\n",
    "history = FullHistory()\n",
    "fg.attach_feature(history)\n",
    "\n",
    "rewrite_graph(fg, include=(\"ShapeOpt\", \"canonicalize\"))\n",
    "\n",
    "# Replay rewrites\n",
    "history.start()\n",
    "pytensor.dprint(fg, print_shape=True)\n",
    "with pytensor.config.change_flags(optimizer_verbose = True):\n",
    "    for i in range(3):\n",
    "        print()\n",
    "        print(\">>> \", end=\"\")\n",
    "        pytensor.dprint(history.next(), print_shape=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "if31pggc4-JF"
   },
   "source": [
    "Does it make sense for this indirect way of rewriting? Indexing is a common operation, we wouldn't want to try to transform the whole graph everytime a user requests a slice of a variable. We do want to avoid doing useless work, so if that's all the user needs, we'll try it. Once we have this machinery in place it's also unnecessary to have a \"transformation\" function for this purpose.\n",
    "\n",
    "This also allows time to reason about the graph holistically. What if we are indexing a variable in one place, but the full output is still needed for another operation? In that case it's better not to optimize the subgraph, as it will result in repeated computations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-11T10:51:56.829649Z",
     "start_time": "2025-03-11T10:51:56.820086Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "BgX4RTJe43Fn",
    "outputId": "d5cad2d7-c976-421a-8175-070499725c50"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Add [id A] shape=()\n",
      " ├─ Subtensor{i, j} [id B] shape=()\n",
      " │  ├─ Blockwise{dot, (m,k),(k,n)->(m,n)} [id C] shape=(512, 256)\n",
      " │  │  ├─ Cos [id D] shape=(512, 64)\n",
      " │  │  │  └─ x [id E] shape=(512, 64)\n",
      " │  │  └─ Exp [id F] shape=(64, 256)\n",
      " │  │     └─ y [id G] shape=(64, 256)\n",
      " │  ├─ 0 [id H] shape=()\n",
      " │  └─ -1 [id I] shape=()\n",
      " └─ Sum{axes=None} [id J] shape=()\n",
      "    └─ Blockwise{dot, (m,k),(k,n)->(m,n)} [id C] shape=(512, 256)\n",
      "       └─ ···\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x7fe53b434d60>"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_out = outs[0, -1] + outs.sum()\n",
    "new_out.dprint(print_shape=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-11T10:51:56.968246Z",
     "start_time": "2025-03-11T10:51:56.957029Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "VoSjBstN5ztS",
    "outputId": "3a37c2f4-ed97-4957-ce44-120955970c2b"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Add [id A] shape=()\n",
      " ├─ Subtensor{i, j} [id B] shape=()\n",
      " │  ├─ dot [id C] shape=(512, 256)\n",
      " │  │  ├─ Cos [id D] shape=(512, 64)\n",
      " │  │  │  └─ x [id E] shape=(512, 64)\n",
      " │  │  └─ Exp [id F] shape=(64, 256)\n",
      " │  │     └─ y [id G] shape=(64, 256)\n",
      " │  ├─ 0 [id H] shape=()\n",
      " │  └─ -1 [id I] shape=()\n",
      " └─ Sum{axes=None} [id J] shape=()\n",
      "    └─ dot [id C] shape=(512, 256)\n",
      "       └─ ···\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x7fe53b434d60>"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rewrite_graph(new_out).dprint(print_shape=True)  # The whole dot is still there!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "h6pACYEXxivg"
   },
   "source": [
    "## Bouns: Shape of a graph\n",
    "\n",
    "It's common to request the shape of graph in PyTensor, and sometimes that's all we actually need, so it's nice if we could get it without having to compute the whole graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-11T10:51:57.053137Z",
     "start_time": "2025-03-11T10:51:57.046585Z"
    },
    "id": "aNWjrykJxhg4"
   },
   "outputs": [],
   "source": [
    "x = pt.matrix(\"x\", shape=(None, None))\n",
    "out = pt.concatenate([x.ravel(), (x @ x.T).ravel()])\n",
    "out_shape = pt.shape(out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-11T10:51:57.119670Z",
     "start_time": "2025-03-11T10:51:57.111222Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "2Nnt_4lhwLE6",
    "outputId": "adb10895-b31e-4ff3-d109-31db2a82a5ca"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape [id A] shape=(1,)\n",
      " └─ Join [id B] shape=(?,)\n",
      "    ├─ 0 [id C] shape=()\n",
      "    ├─ Reshape{1} [id D] shape=(?,)\n",
      "    │  ├─ x [id E] shape=(?, ?)\n",
      "    │  └─ [-1] [id F] shape=(1,)\n",
      "    └─ Reshape{1} [id G] shape=(?,)\n",
      "       ├─ Blockwise{dot, (m,k),(k,n)->(m,n)} [id H] shape=(?, ?)\n",
      "       │  ├─ x [id E] shape=(?, ?)\n",
      "       │  └─ Transpose{axes=[1, 0]} [id I] shape=(?, ?) 'x.T'\n",
      "       │     └─ x [id E] shape=(?, ?)\n",
      "       └─ [-1] [id J] shape=(1,)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x7fe53b434d60>"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out_shape.dprint(print_shape=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-11T10:51:57.172756Z",
     "start_time": "2025-03-11T10:51:57.160073Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "dJhTfs5dy-zC",
    "outputId": "8bee688b-0ede-4b27-9b63-6dfece0e07be"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape [id A] shape=(1,)\n",
      " └─ Join [id B] shape=(10,)\n",
      "    ├─ 0 [id C] shape=()\n",
      "    ├─ Reshape{1} [id D] shape=(6,)\n",
      "    │  ├─ x [id E] shape=(2, 3)\n",
      "    │  └─ [-1] [id F] shape=(1,)\n",
      "    └─ Reshape{1} [id G] shape=(4,)\n",
      "       ├─ Blockwise{dot, (m,k),(k,n)->(m,n)} [id H] shape=(2, 2)\n",
      "       │  ├─ x [id E] shape=(2, 3)\n",
      "       │  └─ Transpose{axes=[1, 0]} [id I] shape=(3, 2)\n",
      "       │     └─ x [id E] shape=(2, 3)\n",
      "       └─ [-1] [id J] shape=(1,)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x7fe53b434d60>"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from pytensor.graph.replace import clone_replace\n",
    "\n",
    "# A bit more readable if we provide a static shape for x, but less interesting\n",
    "static_out_shape = clone_replace(out_shape, {x: pt.matrix(\"x\", shape=(2, 3))})\n",
    "static_out_shape.dprint(print_shape=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "85RgHz2xyL1W"
   },
   "source": [
    "Just like indexing, shape is not really a graph transformation, it's just a symbolic operation that takes as input a variable and returns its shape as the output.\n",
    "\n",
    "Again, the key is in the ability of PyTensor to rewrite and reason about graphs. We we rewrite this graph, PyTensor will be motivated to obtain the most efficient expression of the shape we requested. And that will be a graph -> shape transformation for all intents and purposes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-11T10:51:57.275081Z",
     "start_time": "2025-03-11T10:51:57.222068Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "7-J_G625yExk",
    "outputId": "08fe87bc-9806-4267-86fe-1297efba833c"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MakeVector{dtype='int64'} [id A] shape=(1,)\n",
      " └─ Add [id B] shape=()\n",
      "    ├─ Mul [id C] shape=()\n",
      "    │  ├─ Shape_i{0} [id D] shape=()\n",
      "    │  │  └─ x [id E] shape=(?, ?)\n",
      "    │  └─ Shape_i{1} [id F] shape=()\n",
      "    │     └─ x [id E] shape=(?, ?)\n",
      "    └─ Mul [id G] shape=()\n",
      "       ├─ Shape_i{0} [id D] shape=()\n",
      "       │  └─ ···\n",
      "       └─ Shape_i{0} [id D] shape=()\n",
      "          └─ ···\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x7fe53b434d60>"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Look ma, I can compute the shape without doing any dots, reshapes or joins!\n",
    "rewritten_out_shape = rewrite_graph(out_shape, include=(\"ShapeOpt\", \"canonicalize\"))\n",
    "rewritten_out_shape.dprint(print_shape=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-11T10:51:58.711268Z",
     "start_time": "2025-03-11T10:51:57.301995Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Z6zjvn5KyndH",
    "outputId": "6212e894-28db-45f8-e035-d234378af610"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([10])"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Let's evaluate it\n",
    "rewritten_out_shape.eval({x: np.zeros((2, 3))})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-11T10:51:58.768153Z",
     "start_time": "2025-03-11T10:51:58.751345Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "9Hbqmj6Xzmp5",
    "outputId": "ffd938c1-efa8-4f34-86ee-53fa615d1a77"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[10] [id A] shape=(1,)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x7fe53b434d60>"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# The the boring static case can be constant folded, can't get more efficinet than that\n",
    "rewritten_static_out_shape = rewrite_graph(static_out_shape, include=(\"ShapeOpt\", \"canonicalize\"))\n",
    "rewritten_static_out_shape.dprint(print_shape=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Zyu8nefCz96C"
   },
   "source": [
    "We can think of the shape rewrite machinery as a lazy graph transformation. We ask for some computational property, shape in this case, and let PyTensor reason about it symbolically to arrive at a nice solution.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0ns-8vyt0P5T"
   },
   "source": [
    "### How did it get there?\n",
    "\n",
    "Unlike the indexing example above here we needed to inform PyTensor we were interested in doing shape reasoning with the `ShapeOpt`. It's not part of canonicalization that is run by default.\n",
    "\n",
    "`ShapeOpt` introduces a feature in the FunctionGraph that is being optimized that does a bunch of clever things to reason about shapes. One of the most important is to query each Op that has an `infer_shape` method how it would go about computing the shapes of its outputs if it only knew the shapes of the inputs.\n",
    "\n",
    "Let's take a small peek to understand how PyTensor reasoned about the shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-11T10:51:58.841620Z",
     "start_time": "2025-03-11T10:51:58.810308Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "7Tlv3S4Wz2_e",
    "outputId": "96659f9e-e1a0-498e-d00b-1c3e295a068a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rewriting: rewrite local_shape_to_shape_i replaces Shape.0 of Shape(Join.0) with MakeVector{dtype='int64'}.0 of MakeVector{dtype='int64'}(Switch.0)\n",
      "rewriting: rewrite local_useless_elemwise_comparison replaces Ge.0 of Ge(0, 0) with Second.0 of Second(0, True)\n",
      "rewriting: rewrite local_useless_fill replaces Second.0 of Second(0, True) with True of None\n",
      "rewriting: rewrite local_add_canonizer replaces Add.0 of Add(0, 1) with 1 of None\n",
      "rewriting: rewrite local_useless_switch replaces Switch.0 of Switch(True, Add.0, Mul.0) with Add.0 of Add(Mul.0, Mul.0)\n"
     ]
    }
   ],
   "source": [
    "with pytensor.config.change_flags(optimizer_verbose=True):\n",
    "    rewrite_graph(out_shape, include=(\"ShapeOpt\", \"canonicalize\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-11T10:51:58.928503Z",
     "start_time": "2025-03-11T10:51:58.866689Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "oZpp9XKL0fAP",
    "outputId": "33847d7d-901c-426b-e91b-4d74fe18a6c2"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape [id A] shape=(1,) 5\n",
      " └─ Join [id B] shape=(?,) 4\n",
      "    ├─ 0 [id C] shape=()\n",
      "    ├─ Reshape{1} [id D] shape=(?,) 3\n",
      "    │  ├─ x [id E] shape=(?, ?)\n",
      "    │  └─ [-1] [id F] shape=(1,)\n",
      "    └─ Reshape{1} [id G] shape=(?,) 2\n",
      "       ├─ Blockwise{dot, (m,k),(k,n)->(m,n)} [id H] shape=(?, ?) 1\n",
      "       │  ├─ x [id E] shape=(?, ?)\n",
      "       │  └─ Transpose{axes=[1, 0]} [id I] shape=(?, ?) 'x.T' 0\n",
      "       │     └─ x [id E] shape=(?, ?)\n",
      "       └─ [-1] [id F] shape=(1,)\n",
      "\n",
      ">>> local_shape_to_shape_i\n",
      "MakeVector{dtype='int64'} [id A] shape=(1,) 10\n",
      " └─ Switch [id B] shape=() 9\n",
      "    ├─ Eq [id C] shape=() 8\n",
      "    │  ├─ 0 [id D] shape=()\n",
      "    │  └─ Switch [id E] shape=() 7\n",
      "    │     ├─ Ge [id F] shape=() 6\n",
      "    │     │  ├─ 0 [id G] shape=()\n",
      "    │     │  └─ 0 [id H] shape=()\n",
      "    │     ├─ 0 [id G] shape=()\n",
      "    │     └─ Add [id I] shape=() 5\n",
      "    │        ├─ 0 [id G] shape=()\n",
      "    │        └─ 1 [id J] shape=()\n",
      "    ├─ Add [id K] shape=() 4\n",
      "    │  ├─ Mul [id L] shape=() 2\n",
      "    │  │  ├─ Shape_i{0} [id M] shape=() 1\n",
      "    │  │  │  └─ x [id N] shape=(?, ?)\n",
      "    │  │  └─ Shape_i{1} [id O] shape=() 0\n",
      "    │  │     └─ x [id N] shape=(?, ?)\n",
      "    │  └─ Mul [id P] shape=() 3\n",
      "    │     ├─ Shape_i{0} [id M] shape=() 1\n",
      "    │     │  └─ ···\n",
      "    │     └─ Shape_i{0} [id M] shape=() 1\n",
      "    │        └─ ···\n",
      "    └─ Mul [id L] shape=() 2\n",
      "       └─ ···\n",
      "\n",
      ">>> MergeOptimizer\n",
      "MakeVector{dtype='int64'} [id A] shape=(1,) 10\n",
      " └─ Switch [id B] shape=() 9\n",
      "    ├─ Eq [id C] shape=() 8\n",
      "    │  ├─ 0 [id D] shape=()\n",
      "    │  └─ Switch [id E] shape=() 7\n",
      "    │     ├─ Ge [id F] shape=() 6\n",
      "    │     │  ├─ 0 [id D] shape=()\n",
      "    │     │  └─ 0 [id G] shape=()\n",
      "    │     ├─ 0 [id D] shape=()\n",
      "    │     └─ Add [id H] shape=() 5\n",
      "    │        ├─ 0 [id D] shape=()\n",
      "    │        └─ 1 [id I] shape=()\n",
      "    ├─ Add [id J] shape=() 4\n",
      "    │  ├─ Mul [id K] shape=() 2\n",
      "    │  │  ├─ Shape_i{0} [id L] shape=() 1\n",
      "    │  │  │  └─ x [id M] shape=(?, ?)\n",
      "    │  │  └─ Shape_i{1} [id N] shape=() 0\n",
      "    │  │     └─ x [id M] shape=(?, ?)\n",
      "    │  └─ Mul [id O] shape=() 3\n",
      "    │     ├─ Shape_i{0} [id L] shape=() 1\n",
      "    │     │  └─ ···\n",
      "    │     └─ Shape_i{0} [id L] shape=() 1\n",
      "    │        └─ ···\n",
      "    └─ Mul [id K] shape=() 2\n",
      "       └─ ···\n",
      "\n",
      ">>> MergeOptimizer\n",
      "MakeVector{dtype='int64'} [id A] shape=(1,) 10\n",
      " └─ Switch [id B] shape=() 9\n",
      "    ├─ Eq [id C] shape=() 8\n",
      "    │  ├─ 0 [id D] shape=()\n",
      "    │  └─ Switch [id E] shape=() 7\n",
      "    │     ├─ Ge [id F] shape=() 6\n",
      "    │     │  ├─ 0 [id D] shape=()\n",
      "    │     │  └─ 0 [id D] shape=()\n",
      "    │     ├─ 0 [id D] shape=()\n",
      "    │     └─ Add [id G] shape=() 5\n",
      "    │        ├─ 0 [id D] shape=()\n",
      "    │        └─ 1 [id H] shape=()\n",
      "    ├─ Add [id I] shape=() 4\n",
      "    │  ├─ Mul [id J] shape=() 2\n",
      "    │  │  ├─ Shape_i{0} [id K] shape=() 1\n",
      "    │  │  │  └─ x [id L] shape=(?, ?)\n",
      "    │  │  └─ Shape_i{1} [id M] shape=() 0\n",
      "    │  │     └─ x [id L] shape=(?, ?)\n",
      "    │  └─ Mul [id N] shape=() 3\n",
      "    │     ├─ Shape_i{0} [id K] shape=() 1\n",
      "    │     │  └─ ···\n",
      "    │     └─ Shape_i{0} [id K] shape=() 1\n",
      "    │        └─ ···\n",
      "    └─ Mul [id J] shape=() 2\n",
      "       └─ ···\n",
      "\n",
      ">>> local_useless_elemwise_comparison\n",
      "MakeVector{dtype='int64'} [id A] shape=(1,) 10\n",
      " └─ Switch [id B] shape=() 9\n",
      "    ├─ Eq [id C] shape=() 8\n",
      "    │  ├─ 0 [id D] shape=()\n",
      "    │  └─ Switch [id E] shape=() 7\n",
      "    │     ├─ Second [id F] shape=() 6\n",
      "    │     │  ├─ 0 [id D] shape=()\n",
      "    │     │  └─ True [id G] shape=()\n",
      "    │     ├─ 0 [id D] shape=()\n",
      "    │     └─ Add [id H] shape=() 5\n",
      "    │        ├─ 0 [id D] shape=()\n",
      "    │        └─ 1 [id I] shape=()\n",
      "    ├─ Add [id J] shape=() 4\n",
      "    │  ├─ Mul [id K] shape=() 2\n",
      "    │  │  ├─ Shape_i{0} [id L] shape=() 1\n",
      "    │  │  │  └─ x [id M] shape=(?, ?)\n",
      "    │  │  └─ Shape_i{1} [id N] shape=() 0\n",
      "    │  │     └─ x [id M] shape=(?, ?)\n",
      "    │  └─ Mul [id O] shape=() 3\n",
      "    │     ├─ Shape_i{0} [id L] shape=() 1\n",
      "    │     │  └─ ···\n",
      "    │     └─ Shape_i{0} [id L] shape=() 1\n",
      "    │        └─ ···\n",
      "    └─ Mul [id K] shape=() 2\n",
      "       └─ ···\n",
      "\n",
      ">>> local_useless_fill\n",
      "MakeVector{dtype='int64'} [id A] shape=(1,) 9\n",
      " └─ Switch [id B] shape=() 8\n",
      "    ├─ Eq [id C] shape=() 7\n",
      "    │  ├─ 0 [id D] shape=()\n",
      "    │  └─ Switch [id E] shape=() 6\n",
      "    │     ├─ True [id F] shape=()\n",
      "    │     ├─ 0 [id D] shape=()\n",
      "    │     └─ Add [id G] shape=() 5\n",
      "    │        ├─ 0 [id D] shape=()\n",
      "    │        └─ 1 [id H] shape=()\n",
      "    ├─ Add [id I] shape=() 4\n",
      "    │  ├─ Mul [id J] shape=() 2\n",
      "    │  │  ├─ Shape_i{0} [id K] shape=() 1\n",
      "    │  │  │  └─ x [id L] shape=(?, ?)\n",
      "    │  │  └─ Shape_i{1} [id M] shape=() 0\n",
      "    │  │     └─ x [id L] shape=(?, ?)\n",
      "    │  └─ Mul [id N] shape=() 3\n",
      "    │     ├─ Shape_i{0} [id K] shape=() 1\n",
      "    │     │  └─ ···\n",
      "    │     └─ Shape_i{0} [id K] shape=() 1\n",
      "    │        └─ ···\n",
      "    └─ Mul [id J] shape=() 2\n",
      "       └─ ···\n",
      "\n",
      ">>> local_add_canonizer\n",
      "MakeVector{dtype='int64'} [id A] shape=(1,) 8\n",
      " └─ Switch [id B] shape=() 7\n",
      "    ├─ Eq [id C] shape=() 6\n",
      "    │  ├─ 0 [id D] shape=()\n",
      "    │  └─ Switch [id E] shape=() 5\n",
      "    │     ├─ True [id F] shape=()\n",
      "    │     ├─ 0 [id D] shape=()\n",
      "    │     └─ 1 [id G] shape=()\n",
      "    ├─ Add [id H] shape=() 4\n",
      "    │  ├─ Mul [id I] shape=() 2\n",
      "    │  │  ├─ Shape_i{0} [id J] shape=() 1\n",
      "    │  │  │  └─ x [id K] shape=(?, ?)\n",
      "    │  │  └─ Shape_i{1} [id L] shape=() 0\n",
      "    │  │     └─ x [id K] shape=(?, ?)\n",
      "    │  └─ Mul [id M] shape=() 3\n",
      "    │     ├─ Shape_i{0} [id J] shape=() 1\n",
      "    │     │  └─ ···\n",
      "    │     └─ Shape_i{0} [id J] shape=() 1\n",
      "    │        └─ ···\n",
      "    └─ Mul [id I] shape=() 2\n",
      "       └─ ···\n",
      "\n",
      ">>> constant_folding\n",
      "MakeVector{dtype='int64'} [id A] shape=(1,) 7\n",
      " └─ Switch [id B] shape=() 6\n",
      "    ├─ Eq [id C] shape=() 5\n",
      "    │  ├─ 0 [id D] shape=()\n",
      "    │  └─ 0 [id E] shape=()\n",
      "    ├─ Add [id F] shape=() 4\n",
      "    │  ├─ Mul [id G] shape=() 2\n",
      "    │  │  ├─ Shape_i{0} [id H] shape=() 1\n",
      "    │  │  │  └─ x [id I] shape=(?, ?)\n",
      "    │  │  └─ Shape_i{1} [id J] shape=() 0\n",
      "    │  │     └─ x [id I] shape=(?, ?)\n",
      "    │  └─ Mul [id K] shape=() 3\n",
      "    │     ├─ Shape_i{0} [id H] shape=() 1\n",
      "    │     │  └─ ···\n",
      "    │     └─ Shape_i{0} [id H] shape=() 1\n",
      "    │        └─ ···\n",
      "    └─ Mul [id G] shape=() 2\n",
      "       └─ ···\n",
      "\n",
      ">>> constant_folding\n",
      "MakeVector{dtype='int64'} [id A] shape=(1,) 6\n",
      " └─ Switch [id B] shape=() 5\n",
      "    ├─ True [id C] shape=()\n",
      "    ├─ Add [id D] shape=() 4\n",
      "    │  ├─ Mul [id E] shape=() 2\n",
      "    │  │  ├─ Shape_i{0} [id F] shape=() 1\n",
      "    │  │  │  └─ x [id G] shape=(?, ?)\n",
      "    │  │  └─ Shape_i{1} [id H] shape=() 0\n",
      "    │  │     └─ x [id G] shape=(?, ?)\n",
      "    │  └─ Mul [id I] shape=() 3\n",
      "    │     ├─ Shape_i{0} [id F] shape=() 1\n",
      "    │     │  └─ ···\n",
      "    │     └─ Shape_i{0} [id F] shape=() 1\n",
      "    │        └─ ···\n",
      "    └─ Mul [id E] shape=() 2\n",
      "       └─ ···\n",
      "\n",
      ">>> local_useless_switch\n",
      "MakeVector{dtype='int64'} [id A] shape=(1,) 5\n",
      " └─ Add [id B] shape=() 4\n",
      "    ├─ Mul [id C] shape=() 3\n",
      "    │  ├─ Shape_i{0} [id D] shape=() 0\n",
      "    │  │  └─ x [id E] shape=(?, ?)\n",
      "    │  └─ Shape_i{1} [id F] shape=() 2\n",
      "    │     └─ x [id E] shape=(?, ?)\n",
      "    └─ Mul [id G] shape=() 1\n",
      "       ├─ Shape_i{0} [id D] shape=() 0\n",
      "       │  └─ ···\n",
      "       └─ Shape_i{0} [id D] shape=() 0\n",
      "          └─ ···\n"
     ]
    }
   ],
   "source": [
    "# Now step by step\n",
    "\n",
    "from pytensor.graph import FunctionGraph\n",
    "\n",
    "fg = FunctionGraph(outputs=[out_shape])\n",
    "history = FullHistory()\n",
    "fg.attach_feature(history)\n",
    "\n",
    "rewrite_graph(fg, include=(\"ShapeOpt\", \"canonicalize\"))\n",
    "\n",
    "# Replay rewrites\n",
    "history.start()\n",
    "pytensor.dprint(fg, print_shape=True)\n",
    "with pytensor.config.change_flags(optimizer_verbose = True):\n",
    "    for i in range(9):\n",
    "        print()\n",
    "        print(\">>> \", end=\"\")\n",
    "        pytensor.dprint(history.next(), print_shape=True)\n"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "authorship_tag": "ABX9TyM62pVR72cHnhH99GB8k1xG",
   "include_colab_link": true,
   "provenance": []
  },
  "kernelspec": {
   "display_name": "pytensor-dev",
   "language": "python",
   "name": "pytensor-dev"
  },
  "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}