dynamic_shapes.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "80d04fd5",
   "metadata": {},
   "source": [
    "This notebook implements rank specialized dynamic shapes on [Simple\n",
    "Autograd](https://colab.research.google.com/drive/1VpeE6UvEPRz9HmsHh1KS0XxXjYu533EC?usp=sharing)\n",
    "The goal is to have an easy to hack on prototype of dynamic shapes\n",
    "that you can use to explore different parts of the design space for\n",
    "dynamic shapes.  The dynamic shapes model in this notebook is\n",
    "appropriate for situations when you **truly do not know** the values\n",
    "of dynamic shapes as you are tracing the model (this is the case for\n",
    "backends like XLA, where we are not computing as we execute our\n",
    "program and thus cannot know the shape of data-dependent operations\n",
    "like torch.nonzero without breaking the graph).  This notebook does\n",
    "not describe the case when you are interested in generalizing traces\n",
    "of concrete executions.\n",
    "\n",
    "Most of the simplest graph capture mechanisms require shape\n",
    "specialization, because they simply proceed by running an actual\n",
    "iteration of the computation with real inputs, and simply recording\n",
    "everything that occurred during the process.  This causes problems,\n",
    "however, when shapes vary across different runs (e.g., you have a\n",
    "dynamically sized input or you use data-dependent operators like\n",
    "nonzero/unique).  So logically, you'd like some way to record \"hey,\n",
    "this shape isn't 1024, it can vary, don't make assumptions based on it\n",
    "happening to be 1024 this time.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "18925177",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch._C.Generator at 0x107453770>"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import functools\n",
    "import itertools\n",
    "import traceback\n",
    "from dataclasses import dataclass, field\n",
    "from enum import Enum\n",
    "from typing import Any, Callable, Dict, List, NamedTuple, Optional, Tuple, Union\n",
    "\n",
    "import torch\n",
    "\n",
    "torch.manual_seed(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5f8e95f",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "To get started, we define some utility functions the autograd\n",
    "implementation.  FreshSupply lets us generate fresh names for\n",
    "variables and symbolic integers in our program (for clarity,\n",
    "we give them separate prefixes: v for variable, i for integer).\n",
    "`gradient_tape` contains the global autograd tape for our\n",
    "program; for more details, read the exposition in [Simple\n",
    "Autograd](https://colab.research.google.com/drive/1VpeE6UvEPRz9HmsHh1KS0XxXjYu533EC?usp=sharing)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f823fde8",
   "metadata": {
    "lines_to_end_of_cell_marker": 2
   },
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class FreshSupply:\n",
    "    prefix: str\n",
    "    fresh: int = 0\n",
    "\n",
    "    def __call__(self):\n",
    "        r = f\"{self.prefix}{self.fresh}\"\n",
    "        self.fresh += 1\n",
    "        return r\n",
    "\n",
    "\n",
    "fresh_var = FreshSupply(\"v\")\n",
    "fresh_int = FreshSupply(\"i\")\n",
    "\n",
    "\n",
    "class TapeEntry(NamedTuple):\n",
    "    # names of the inputs to the original computation\n",
    "    inputs: List[str]\n",
    "    # names of the outputs of the original computation\n",
    "    outputs: List[str]\n",
    "    # apply chain rule\n",
    "    propagate: \"Callable[List[Variable], List[Variable]]\"\n",
    "\n",
    "\n",
    "gradient_tape: List[TapeEntry] = []\n",
    "\n",
    "\n",
    "def reset():\n",
    "    gradient_tape.clear()\n",
    "    fresh_var.fresh = 0\n",
    "    fresh_int.fresh = 0\n",
    "    CURRENT_GRAPH.nodes.clear()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f355524f",
   "metadata": {},
   "source": [
    "The raison d'etre of dynamic shapes is for graph capture, and so to\n",
    "start we will define a little IR representing sequences of operations\n",
    "that we will be recording.  To avoid string confusion, we'll use\n",
    "a little wrapper class Op to represent each operation; we'll give\n",
    "meaning to these objects later by defining interpretation rules for\n",
    "them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9cef476c",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "@dataclass(frozen=True)\n",
    "class Op:\n",
    "    name: str\n",
    "\n",
    "    def __str__(self):\n",
    "        return self.name\n",
    "\n",
    "\n",
    "# In this example, we only have asserts over integers, but you could\n",
    "# imagine other integer operations, e.g., addition, for handling\n",
    "# operators like torch.cat\n",
    "int_assert_eq = Op(\"int_assert_eq\")\n",
    "var_constant = Op(\"var_constant\")\n",
    "var_add = Op(\"var_add\")\n",
    "var_mul = Op(\"var_mul\")\n",
    "var_sum = Op(\"var_sum\")\n",
    "var_expand = Op(\"var_expand\")\n",
    "# Unlike nonzero, this returns both the number of nonzero elements, as\n",
    "# well as the return tensor\n",
    "var_nonzero_impl = Op(\"var_nonzero_impl\")\n",
    "var_index = Op(\"var_index\")  # aka x[i]\n",
    "var_index_backward = Op(\"var_index_backward\")\n",
    "var_squeeze = Op(\"squeeze\")\n",
    "var_unsqueeze = Op(\"unsqueeze\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a77b3abb",
   "metadata": {},
   "source": [
    "Let's write a little interpreter for these these operations.  Their\n",
    "implementations will proceed as you might expect.  Since there are a\n",
    "lot of operators, we'll put all of the functions for handling these\n",
    "operations in an INTERP_RULES dictionary.  The interpreter ops\n",
    "distinguish between args and kwargs; args are dynamic data that is\n",
    "allowed to depend on other operations in our IR, whereas kwargs\n",
    "is static data which never depends on other computation.\n",
    "\n",
    "Note that operations on integers are also represented in the graph!\n",
    "In some graph representations, integers are represented as scalar\n",
    "tensors, but for clarity in this presentation they are represented by\n",
    "a separate type."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "73b05418",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "def register(d, k):\n",
    "    def inner(f):\n",
    "        d[k] = f\n",
    "\n",
    "    return inner\n",
    "\n",
    "\n",
    "INTERP_RULES = {}\n",
    "\n",
    "\n",
    "@register(INTERP_RULES, int_assert_eq)\n",
    "def interp_int_assert_eq(x: int, y: int):\n",
    "    assert x == y\n",
    "\n",
    "\n",
    "# unlike nonzero, this also returns the symbolic shape\n",
    "# because this is an \"existential telescope\" the fresh shape gotta come\n",
    "# first\n",
    "@register(INTERP_RULES, var_nonzero_impl)\n",
    "def interp_var_nonzero_impl(x: torch.Tensor):\n",
    "    r = torch.nonzero(x)\n",
    "    return r.shape[0], r\n",
    "\n",
    "\n",
    "INTERP_RULES[var_index] = lambda t, i: t[i]\n",
    "# NB: this is inefficient: t's data doesn't to be retained to allocate\n",
    "# the zeros, only the dtype (actually technically inferrable from g) and size\n",
    "INTERP_RULES[var_index_backward] = lambda t, i, g: torch.zeros_like(t).index_put(\n",
    "    (i,), g, accumulate=True\n",
    ")\n",
    "INTERP_RULES[var_constant] = lambda *, val: val\n",
    "INTERP_RULES[var_add] = lambda x, y: x + y\n",
    "INTERP_RULES[var_mul] = lambda x, y: x * y\n",
    "INTERP_RULES[var_sum] = lambda x: x.sum()\n",
    "INTERP_RULES[var_expand] = lambda x, sizes: x.expand(sizes)\n",
    "INTERP_RULES[var_squeeze] = lambda x, *, dim: x.squeeze(dim)\n",
    "INTERP_RULES[var_unsqueeze] = lambda x, *, dim: x.unsqueeze(dim)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f334b034",
   "metadata": {},
   "source": [
    "Let's actually define IR nodes that make use of these operations.\n",
    "In many IRs, you only allow variables as arguments (so called\n",
    "\"administrative normal form\"); but in this treatment, I chose to also\n",
    "allow tuples in atoms to make the IR easier to read (as we can just\n",
    "directly specify shapes as arguments, rather than having to first\n",
    "allocate a tuple and then pass it the function).  There is no\n",
    "need for first class tuples because we are dealing with *rank\n",
    "specialized* dynamic shapes, so the size of tuples is always known\n",
    "ahead of time.  Similarly, if an integer argument in a shape is\n",
    "known statically ahead of time, we'll let you just inline it into\n",
    "the argument site.  We will call these structures atoms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "212bcc97",
   "metadata": {
    "lines_to_end_of_cell_marker": 2
   },
   "outputs": [],
   "source": [
    "Atom = Union[str, int, List[Union[str, int]]]\n",
    "\n",
    "\n",
    "def str_atom(a: Atom) -> str:\n",
    "    if isinstance(a, str):\n",
    "        return a\n",
    "    elif isinstance(a, int):\n",
    "        return str(a)\n",
    "    else:\n",
    "        return f\"({', '.join(str_atom(b) for b in a)})\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d45e391",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "Our node looks fairly similar to an FX node: given some list of input\n",
    "atoms and a dictionary of arbitrary extra (static) parameters, run\n",
    "an operator on it and bind the results to outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f05429d8",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class Node:\n",
    "    op: Op\n",
    "    inputs: List[Atom]\n",
    "    outputs: List[str]\n",
    "    params: Dict[str, Any] = field(default_factory=dict)\n",
    "\n",
    "    def __str__(self):\n",
    "        outputs_str = \", \".join(self.outputs)\n",
    "        outputs_str += \" = \" if self.outputs else \"\"\n",
    "        inputs_str = \", \".join(str_atom(a) for a in self.inputs)\n",
    "        params_str = \", \" if self.inputs and self.params else \"\"\n",
    "        params_str += \", \".join(f\"{k}={v}\" for k, v in self.params.items())\n",
    "        return f\"{outputs_str}{self.op}({inputs_str}{params_str})\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37f886c4",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "A graph is simply a list of nodes.  In a real application we may also\n",
    "need some provision for graph-level inputs/outputs, but they aren't\n",
    "necessary for these example so they are omitted."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "66891836",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class Graph:\n",
    "    nodes: List[Node] = field(default_factory=list)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70dc1d92",
   "metadata": {},
   "source": [
    "Now, we can finish our interpreter on graphs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5c45f3aa",
   "metadata": {
    "lines_to_end_of_cell_marker": 2
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "def tuplify(outs):\n",
    "    if outs is None:\n",
    "        return ()\n",
    "    elif isinstance(outs, tuple):\n",
    "        return outs\n",
    "    else:\n",
    "        return (outs,)\n",
    "\n",
    "\n",
    "def interp_atom(atom: Atom, env: Dict[str, Any]):\n",
    "    if isinstance(atom, str):\n",
    "        return env[atom]\n",
    "    elif isinstance(atom, tuple):\n",
    "        return tuple(interp_atom(a, env) for a in atom)\n",
    "    else:\n",
    "        return atom\n",
    "\n",
    "\n",
    "def interp_inputs(inputs: List[Atom], env: Dict[str, Any]):\n",
    "    return tuple(interp_atom(i, env) for i in inputs)\n",
    "\n",
    "\n",
    "# Mutates the environment, storing the results into env\n",
    "def interp_node(n: Node, env: Dict[str, Any]):\n",
    "    args = tuple(interp_atom(i, env) for i in n.inputs)\n",
    "    try:\n",
    "        outs = tuplify(INTERP_RULES[n.op](*args, **n.params))\n",
    "    except Exception:\n",
    "        print(f\"Failed during: {n}\")\n",
    "        raise\n",
    "    assert len(outs) == len(n.outputs)\n",
    "    for k, v in zip(n.outputs, outs):\n",
    "        env[k] = v"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "293be79f",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "To put it all together, we have a final `interp_graph` function,\n",
    "which takes in initial values for variables and then runs the\n",
    "graph recorded at `CURRENT_GRAPH` to compute some output variables\n",
    "(which it then prints.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "e02eed04",
   "metadata": {},
   "outputs": [],
   "source": [
    "def interp_graph(init: Dict[Union[\"Variable\", \"SymbolicIntNode\"], Any], **outs):\n",
    "    env = {k.name: v for k, v in init.items()}\n",
    "    for n in CURRENT_GRAPH.nodes:\n",
    "        interp_node(n, env)\n",
    "    for k, v in outs.items():\n",
    "        print(f\"{k} = {env[v.name]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f764ff38",
   "metadata": {},
   "source": [
    "Phew, that's it for the IR representation.  Let's actually generate\n",
    "some graphs!  We'll maintain a global current graph which we record\n",
    "into."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "276f43ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "CURRENT_GRAPH = Graph()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c05b151",
   "metadata": {},
   "source": [
    "Like FX, XLA and LazyTensor, we will generate graphs by maintaining\n",
    "symbolic \"proxy\" objects (Variable, SymbolicIntNode) which act\n",
    "like normal tensors/integers but don't actually contain any data\n",
    "and record the operations that occur on them.  Our symbolic integers\n",
    "in this example have a very impoverished interface; in fact they\n",
    "support no operations at all, they simply have a name corresponding\n",
    "to their name in the graph.  (In a tracing JIT context, we might\n",
    "secretly know what the actual values of our symbolic ints are, but\n",
    "for conceptual clarity we do not allow ourselves this in this\n",
    "notebook.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d245f30b",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "class SymbolicIntNode:\n",
    "    name: str\n",
    "\n",
    "    def __init__(self, name=None):\n",
    "        self.name = name or fresh_int()\n",
    "\n",
    "    def __repr__(self):\n",
    "        return self.name\n",
    "\n",
    "\n",
    "class Variable:\n",
    "    shape: List[Union[\"SymInt\", int]]\n",
    "    name: str\n",
    "    dtype: torch.dtype\n",
    "\n",
    "    def __init__(self, shape, dtype: torch.dtype, *, name: str = None):\n",
    "        self.shape = shape\n",
    "        self.dtype = dtype\n",
    "        self.name = name or fresh_var()\n",
    "\n",
    "    def dim(self):\n",
    "        return len(self.shape)\n",
    "\n",
    "    # We need to start with some tensors whose values were not computed\n",
    "    # inside the autograd. This function constructs leaf nodes.\n",
    "    @staticmethod\n",
    "    def constant(value: torch.Tensor, name: str = None):\n",
    "        return record_var(var_constant, tuple(value.shape), value.dtype, val=value)\n",
    "\n",
    "    def __repr__(self):\n",
    "        return f\"{self.name}: {self.shape}\"\n",
    "\n",
    "    # Most of the actual implementations of these operators will be\n",
    "    # given later.\n",
    "\n",
    "    def __mul__(self, rhs: \"Variable\") -> \"Variable\":\n",
    "        # defined later in the notebook\n",
    "        return operator_mul(self, rhs)\n",
    "\n",
    "    def __add__(self, rhs: \"Variable\") -> \"Variable\":\n",
    "        return operator_add(self, rhs)\n",
    "\n",
    "    def sum(self, name: Optional[str] = None) -> \"Variable\":\n",
    "        return operator_sum(self, name)\n",
    "\n",
    "    def expand(self, sizes: List[\"SymInt\"]) -> \"Variable\":\n",
    "        return operator_expand(self, sizes)\n",
    "\n",
    "    def nonzero(self) -> \"Variable\":\n",
    "        return operator_nonzero(self)\n",
    "\n",
    "    def squeeze(self, dim: int) -> \"Variable\":\n",
    "        return operator_squeeze(self, dim)\n",
    "\n",
    "    def unsqueeze(self, dim: int) -> \"Variable\":\n",
    "        return operator_unsqueeze(self, dim)\n",
    "\n",
    "    def zeros_like(self) -> \"Variable\":\n",
    "        return zeros_like(self)\n",
    "\n",
    "    def __getitem__(self, index) -> \"Variable\":\n",
    "        return operator_index(self, index)\n",
    "\n",
    "    def index_backward(self, index, grad_output) -> \"Variable\":\n",
    "        return operator_index_backward(self, index, grad_output)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "304d006d",
   "metadata": {},
   "source": [
    "We allow integers to be specialized (indeed, in lazy tensor, most\n",
    "integers in networks will be statically known), so you may be\n",
    "dealing with a literal integer or a symbolic integer node.  SymInt\n",
    "captures these two possibilities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "9210de39",
   "metadata": {},
   "outputs": [],
   "source": [
    "SymInt = Union[SymbolicIntNode, int]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c42d68ae",
   "metadata": {},
   "source": [
    "To record a graph node involving integers (e.g., `int_assert_eq`),\n",
    "we simply take in the input Variables/SymInts, extract out their\n",
    "names, and record an operator to the current graph for the operation\n",
    "we tried to do.  Furthermore, if the operation returns a result\n",
    "(e.g., an int or a tensor), we have to construct a new proxy object\n",
    "representing the result.  We provide helper functions for doing this\n",
    "with operations that return nothing, a single int, or a single\n",
    "variable (note that the to return a new variable proxy, we must\n",
    "provide a (possibly symbolic but definitely rank specialized) shape\n",
    "and a (concrete) dtype of the result.  (dtypes are concrete because\n",
    "we don't support symbolic dtype!  That is a much rarer thing to want\n",
    "to do inside of a network.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "dfcec677",
   "metadata": {
    "lines_to_end_of_cell_marker": 2
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "def record_arg(a):\n",
    "    if isinstance(a, tuple):\n",
    "        return tuple(record_arg(b) for b in a)\n",
    "    elif isinstance(a, int):\n",
    "        return a\n",
    "    else:\n",
    "        assert isinstance(a, (Variable, SymbolicIntNode))\n",
    "        return a.name\n",
    "\n",
    "\n",
    "def record_none(op, *args, **kwargs):\n",
    "    n = Node(op, tuple(record_arg(a) for a in args), [], kwargs)\n",
    "    print(n)\n",
    "    CURRENT_GRAPH.nodes.append(n)\n",
    "\n",
    "\n",
    "def record_int(op, *args, **kwargs):\n",
    "    i = SymbolicIntNode()\n",
    "    n = Node(op, tuple(a.name for a in args), [i.name], kwargs)\n",
    "    print(n)\n",
    "    CURRENT_GRAPH.nodes.append(n)\n",
    "    return i\n",
    "\n",
    "\n",
    "def record_var(\n",
    "    op, shape: Tuple[SymInt, ...], dtype: torch.dtype, *args, name=None, **kwargs\n",
    "):\n",
    "    r = Variable(shape, dtype, name=name)\n",
    "    n = Node(op, tuple(record_arg(a) for a in args), [r.name], kwargs)\n",
    "    print(f\"{n} : {r.shape}\")\n",
    "    CURRENT_GRAPH.nodes.append(n)\n",
    "    return r"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6afb1172",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "We our now ready to define our first operation: `int_assert_eq`.\n",
    "Although we peephole a few cases where this assertion is vacuously\n",
    "true, in general we don't know if two symbolic integers are actually\n",
    "equal or not, so we just record an assertion in the graph to be\n",
    "verified later at runtime."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "ebfbf3ce",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "def assert_int_eq(x: SymInt, y: SymInt):\n",
    "    # peephole optimization\n",
    "    if (\n",
    "        isinstance(x, SymbolicIntNode)\n",
    "        and isinstance(y, SymbolicIntNode)\n",
    "        and x.name == y.name\n",
    "    ):\n",
    "        return\n",
    "    if isinstance(x, int) and isinstance(y, int) and x == y:\n",
    "        return\n",
    "    record_none(int_assert_eq, x, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "823da943",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "It is not implemented in this notebook, but what if you wanted to\n",
    "immediately check this assertion *while tracing* (instead of waiting\n",
    "for concrete inputs/sizes).  There is now an interesting choice you have\n",
    "to make.  In particular, suppose x is a Tensor whose shape is\n",
    "symbolic; what does this program mean?\n",
    "\n",
    "```\n",
    "assert x.shape[0] == 4\n",
    "```\n",
    "\n",
    "Should you error during tracing or not?  There are two possibilities:\n",
    "\n",
    "1. **The program is invalid.**  In this interpretation, the\n",
    "   shape of x being symbolic is a claim that this program should\n",
    "   work for any choice of x, and now we are trying to **check** that\n",
    "   this is actually true.  When this assert occurs, we say, \"The user\n",
    "   told us x.shape[0] could be anything, but now we see that if\n",
    "   x.shape[0] is not 4 this assert will fail; there's a bug and\n",
    "   we should report an error.\"  This is the sort of thing you would\n",
    "   do if you were given a model with specific sizes annotated as\n",
    "   \"dynamic\" and then were trying to trace the model under this\n",
    "   assumption.  We call these symbolic sizes **rigid**, because\n",
    "   they never change as we execute a program.\n",
    "\n",
    "   Note: it's not strictly required to let a size range over all\n",
    "   integers; for example, a user could specify preconditions for their\n",
    "   model input shapes (e.g., \"x.shape[0] must be even\") which could\n",
    "   then be used to show that asserts on those shapes must be true.  In\n",
    "   the degenerate case, the preconditions specify concrete values for\n",
    "   every shape in the program, making this equivalent to the concrete\n",
    "   shapes case.  In XLA, the information about sizes having upper\n",
    "   bounds serves a similar role.\n",
    "\n",
    "2. **The program is valid.**  In this interpretation, the shape\n",
    "   of x being symbolic just means we don't know anything about the\n",
    "   preconditions/postconditions of our program, and we are trying to\n",
    "   **infer** them from program.  When this assert occurs, we say, \"Aha!\n",
    "   Now we know that x.shape[0] is 4\" and can use that fact in later\n",
    "   analysis (e.g., to find a contradiction with a later assert, which\n",
    "   would indicate that there is a bug in the program).  This is the\n",
    "   sort of thing you would do if you were given a model with no\n",
    "   input/output shape annotations and were trying to infer what the top\n",
    "   level annotations should be.  We call these symbolic sizes\n",
    "   **flexible**, because their nature will change based on the program\n",
    "   we run them on.\n",
    "\n",
    "If you told me I could only support one case, I would pick case (1).\n",
    "To see why, consider torch.nonzero(), a function whose output size is\n",
    "dependent on the data inside the function.  If we wish to write down\n",
    "the shape of this function without reference to the data in the tensor\n",
    "(which is typically what we want to do--we usually want to write the\n",
    "shapes of our programs in a data oblivious way), all we can really say\n",
    "is that there *exists* some (symbolic) size such that the tensor has\n",
    "that size, but no, I can't tell you what it is.  If the user then\n",
    "passed this result tensor into an operator that expects the size to\n",
    "actually be four, we would expect this to be an error.  (Now, it's\n",
    "possible that there *accidentally* were four nonzero elements, but\n",
    "I wouldn't bet on it!)\n",
    "\n",
    "Case (2) has some useful applications; however.  If you are given an\n",
    "arbitrary model with no annotations, you can replace all of the input\n",
    "sizes with flexible sizes, run the model, and get symbolic formulas\n",
    "for what the implicit preconditions for the model are, as well as\n",
    "direct formulas for the output sizes.  If you have a spot in the\n",
    "middle of your model where you're not sure what the size should be,\n",
    "you could stick a flexible size there and have the symbolic algebra\n",
    "system tell you what the size has to be (LazyModule style).\n",
    "\n",
    "In fact, case (1) and case (2) can coexist in the same model.  There\n",
    "are a few ways to do it: you could distinguish between rigid/flexible\n",
    "symbolic variables and have asserts work differently depending on the\n",
    "variable in question/you could also introduce two different types\n",
    "of asserts (one which says \"this assert is derivable from the\n",
    "shape preconditions of this model\" and another which says \"this assert\n",
    "represents external knowledge that I have about the model, and you can\n",
    "now use this for further reasoning.\")  They're not implemented here,\n",
    "but it is a good exercise to try."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74fe8b69",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "We our now ready to implement pointwise multiplication.  In this\n",
    "notebook, we have chosen to faithfully replicate *broadcasting*\n",
    "semantics, which means that adding two tensors with two different\n",
    "sizes at a dimension is OK if one of the sizes is one.\n",
    "\n",
    "To do this, we need a way of taking a (possibly) symbolic integer\n",
    "and checking (at trace time) if it is one.  For dynamic sizes,\n",
    "there is no way to generically do this, because we haven't run\n",
    "the program, we have no idea what the size will be!  So the simplest\n",
    "possible choice is to only report a SymInt as one if it is\n",
    "*specialized* to be one (e.g., it is literally a one integer.)\n",
    "\n",
    "It's important to emphasize that this constitutes a UI divergence\n",
    "from a traditional PyTorch program that tests if a size is 1; if\n",
    "the shape is dynamic, we will report False, but if the shape HAPPENS\n",
    "to later to be instantiated to be 1, an actual PyTorch program would\n",
    "have reported True (whereas we have baked in trace behavior for it\n",
    "being False).  This is another reason why it is pedagogically clearer\n",
    "to avoid having real values \"in your back pocket\", because with them\n",
    "the temptation is to just peek at the actual value to do the test\n",
    "(and this results in a bunch of extra obligations to deal with.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "148ca4ac",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "def definitely_one(x: SymInt) -> bool:\n",
    "    return isinstance(x, int) and x == 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08c2ee4b",
   "metadata": {},
   "source": [
    "This reasoning is not complete.  Suppose x has size (s0,) and y has\n",
    "size (s1,), and we have:\n",
    "\n",
    "```\n",
    "assert x.shape[0] == 1\n",
    "return x + y\n",
    "```\n",
    "\n",
    "Does broadcasting occur on this addition?  A user might reasonably\n",
    "expect that after this assertion, surely broadcasting should occur.\n",
    "But if we are relying on x.shape[0] to *literally* be one (as we\n",
    "are in the implementation of `definitely_one` above), we won't figure\n",
    "it out!  It's not too difficult to figure it out though: in\n",
    "particular, if you implemented unification that would suffice.\n",
    "\n",
    "The trouble is, it's not well specified *how much* symbolic reasoning\n",
    "we should be willing to do.  It seems that unification is necessary,\n",
    "lest obvious instances of transitivity don't work (x == 1, y == x, z\n",
    "== y, operations on z should broadcast).  On the flip side, we\n",
    "shouldn't be in the business of proving arbitrary mathematical\n",
    "theorems (or even shelling to Z3) to figure out if a value is always\n",
    "one in some context.  But what about a ResNet style architecture,\n",
    "where the output size of layers gets reduced and reduced until it\n",
    "hits 1; should we be able to infer that the output of a ResNet in this\n",
    "case is size 1 and eligible for broadcasting?  (Does this even matter,\n",
    "since the input sizes in such networks are typically static and we\n",
    "wouldn't have a dynamic shape in this case anyway?)  To definitively\n",
    "answer these questions would require an analysis of broadcasting usage\n",
    "in the wild (or perhaps an analysis of networks with dynamic sizes.)\n",
    "\n",
    "There is sense in which unification is \"good enough\", however.  If our\n",
    "symbolic reasoning is insufficient for a user, they can always add an\n",
    "assert that a shape in question is one to force broadcasting to occur\n",
    "in that case.  Then, the problem reduces to letting a user know\n",
    "that this is what they ought to do; if two tensors don't broadcast\n",
    "with each other, we may just require their sizes to be the same; but\n",
    "it might not be obvious (without the help of a solver like Z3) that\n",
    "two symbolic sizes are different.  If we lowered to an IR with\n",
    "non-broadcasting operations, this will manifest at runtime where we'll\n",
    "say \"Couldn't add tensors with sizes 1 and N\" (even though the surface\n",
    "language supported broadcasting).  So you want to help the user out\n",
    "here with a better error message, in that case."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a80077cf",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "Anyway, now that we have `definitely_one`, we can implement an\n",
    "assertion that two shapes are broadcastable (returning the\n",
    "possible symbolic shape out the end).  This broadcasting is more\n",
    "conservative than actual PyTorch, because if a shape is dynamic\n",
    "it will reject it (even if at runtime it happened to be one and\n",
    "therefore the broacasting woudl have succeeded)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "3de3a0a2",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "def assert_shape_broadcast(lhs, rhs):\n",
    "    r = []\n",
    "    for x, y in itertools.zip_longest(\n",
    "        reversed(lhs.shape), reversed(rhs.shape), fillvalue=1\n",
    "    ):\n",
    "        if definitely_one(x):\n",
    "            r.append(y)\n",
    "        elif definitely_one(y):\n",
    "            r.append(x)\n",
    "        else:\n",
    "            assert_int_eq(x, y)\n",
    "            r.append(x)  # pick one arbitrarily\n",
    "    return tuple(reversed(r))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55fe07da",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "Finally, we can implement multplication!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "bb72ec81",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "def operator_mul(self: Variable, rhs: Variable) -> Variable:\n",
    "    if isinstance(rhs, float) and rhs == 1.0:\n",
    "        # peephole optimization\n",
    "        return self\n",
    "\n",
    "    # Broadcast the two shapes together, getting the result shape\n",
    "    shape = assert_shape_broadcast(self, rhs)\n",
    "    # We didn't implement type promotion, so just assert that the\n",
    "    # inputs are the same dtype.\n",
    "    assert self.dtype == rhs.dtype\n",
    "    # Record the operation into the graph\n",
    "    r = record_var(var_mul, shape, self.dtype, self, rhs)\n",
    "\n",
    "    # Record what the inputs and outputs of the op were\n",
    "    inputs = [self.name, rhs.name]\n",
    "    outputs = [r.name]\n",
    "\n",
    "    # Define backprop.  This closes over self and rhs, which are\n",
    "    # necessary to define the backwards rule.\n",
    "    def propagate(dL_doutputs: List[Variable]):\n",
    "        (dL_dr,) = dL_doutputs\n",
    "\n",
    "        dr_dself = rhs  # partial derivative of r = self*rhs\n",
    "        dr_drhs = self  # partial derivative of r = self*rhs\n",
    "\n",
    "        # chain rule propagation from outputs to inputs of multiply\n",
    "        dL_dself = dL_dr * dr_dself\n",
    "        dL_drhs = dL_dr * dr_drhs\n",
    "        dL_dinputs = [dL_dself, dL_drhs]\n",
    "        return dL_dinputs\n",
    "\n",
    "    # Finally, we record the compute we did on the tape\n",
    "    gradient_tape.append(TapeEntry(inputs=inputs, outputs=outputs, propagate=propagate))\n",
    "    return r"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c55534f7",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "The implementation of gradient computation works exactly the same\n",
    "way it did in Simple Grad."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "865affed",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "def grad(L, desired_results: List[Variable]) -> List[Variable]:\n",
    "    # this map holds dL/dX for all values X\n",
    "    dL_d: Dict[str, Variable] = {}\n",
    "    # It starts by initializing the 'seed' dL/dL, which is 1\n",
    "    dL_d[L.name] = Variable.constant(torch.ones(()))\n",
    "    print(f\"d{L.name} ------------------------\")\n",
    "\n",
    "    # look up dL_dentries. If a variable is never used to compute the loss,\n",
    "    # we consider its gradient None, see the note below about zeros for more information.\n",
    "    def gather_grad(entries: List[str]):\n",
    "        return [dL_d[entry] if entry in dL_d else None for entry in entries]\n",
    "\n",
    "    # propagate the gradient information backward\n",
    "    for entry in reversed(gradient_tape):\n",
    "        dL_doutputs = gather_grad(entry.outputs)\n",
    "        if all(dL_doutput is None for dL_doutput in dL_doutputs):\n",
    "            # optimize for the case where some gradient pathways are zero. See\n",
    "            # The note below for more details.\n",
    "            continue\n",
    "\n",
    "        # perform chain rule propagation specific to each compute\n",
    "        dL_dinputs = entry.propagate(dL_doutputs)\n",
    "\n",
    "        # Accululate the gradient produced for each input.\n",
    "        # Each use of a variable produces some gradient dL_dinput for that\n",
    "        # use. The multivariate chain rule tells us it is safe to sum\n",
    "        # all the contributions together.\n",
    "        for input, dL_dinput in zip(entry.inputs, dL_dinputs):\n",
    "            if input not in dL_d:\n",
    "                dL_d[input] = dL_dinput\n",
    "            else:\n",
    "                dL_d[input] = dL_d[input] + dL_dinput\n",
    "\n",
    "    # print some information to understand the values of each intermediate\n",
    "    for name, value in dL_d.items():\n",
    "        print(f\"d{L.name}_d{name}: {value.shape} = {value.name}\")\n",
    "    print(f\"------------------------\")\n",
    "\n",
    "    return gather_grad(desired.name for desired in desired_results)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91153e63",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "Add, sum and expand look identical to their versions in Simple Grad.\n",
    "One thing to note, however: expand takes sizes as input, and those\n",
    "sizes can be symbolic!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "75a9e9e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def operator_add(self: Variable, rhs: Variable) -> Variable:\n",
    "    shape = assert_shape_broadcast(self, rhs)\n",
    "    assert self.dtype == rhs.dtype  # no type promotion\n",
    "    r = record_var(var_add, shape, self.dtype, self, rhs)\n",
    "\n",
    "    def propagate(dL_doutputs: List[Variable]):\n",
    "        (dL_dr,) = dL_doutputs\n",
    "        dr_dself = 1.0\n",
    "        dr_drhs = 1.0\n",
    "        dL_dself = dL_dr * dr_dself\n",
    "        dL_drhs = dL_dr * dr_drhs\n",
    "        return [dL_dself, dL_drhs]\n",
    "\n",
    "    gradient_tape.append(\n",
    "        TapeEntry(inputs=[self.name, rhs.name], outputs=[r.name], propagate=propagate)\n",
    "    )\n",
    "    return r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "e738c7d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def operator_sum(self: Variable, name: Optional[str]) -> \"Variable\":\n",
    "    r = record_var(var_sum, (), self.dtype, self, name=name)\n",
    "\n",
    "    def propagate(dL_doutputs: List[Variable]):\n",
    "        (dL_dr,) = dL_doutputs\n",
    "        size = self.shape\n",
    "        return [dL_dr.expand(size)]\n",
    "\n",
    "    gradient_tape.append(\n",
    "        TapeEntry(inputs=[self.name], outputs=[r.name], propagate=propagate)\n",
    "    )\n",
    "    return r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "95390cb8",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "# NB: -1 sizes was not implemented\n",
    "def operator_expand(self: Variable, sizes: List[SymInt]) -> \"Variable\":\n",
    "    assert self.dim() == 0  # only works for scalars\n",
    "    r = record_var(var_expand, sizes, self.dtype, self, sizes)\n",
    "\n",
    "    def propagate(dL_doutputs: List[Variable]):\n",
    "        (dL_dr,) = dL_doutputs\n",
    "        return [dL_dr.sum()]\n",
    "\n",
    "    gradient_tape.append(\n",
    "        TapeEntry(inputs=[self.name], outputs=[r.name], propagate=propagate)\n",
    "    )\n",
    "    return r"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ac8b286",
   "metadata": {},
   "source": [
    "With these operators, we can reprise the simple add-multiply example\n",
    "from the original Simple Grad, but this time first symbolically\n",
    "tracing the graph, and then executing it after the fact."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "ee39ef2b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "v0 = var_add(a, b) : (4,)\n",
      "v1 = var_mul(v0, b) : (4,)\n",
      "v2 = var_constant(val=1.0) : ()\n",
      "dv1 ------------------------\n",
      "v3 = var_mul(v2, b) : (4,)\n",
      "v4 = var_mul(v2, v0) : (4,)\n",
      "v5 = var_add(v4, v3) : (4,)\n",
      "dv1_dv1: () = v2\n",
      "dv1_dv0: (4,) = v3\n",
      "dv1_db: (4,) = v5\n",
      "dv1_da: (4,) = v3\n",
      "------------------------\n",
      "da = tensor([-1.0845, -1.3986,  0.4033,  0.8380])\n",
      "db = tensor([-0.6280, -3.0906, -1.3721,  2.2445])\n"
     ]
    }
   ],
   "source": [
    "\n",
    "\n",
    "def simple(a, b):\n",
    "    t = a + b\n",
    "    return t * b\n",
    "\n",
    "\n",
    "reset()  # reset any compute from other cells\n",
    "# We do no computation in this part, we're just tracing!\n",
    "a = Variable((4,), dtype=torch.float, name=\"a\")\n",
    "b = Variable((4,), dtype=torch.float, name=\"b\")\n",
    "loss = simple(a, b)\n",
    "da, db = grad(loss, [a, b])\n",
    "# Setting a and b to random tensors, run the interpreted graph\n",
    "# and print out the result of da and db.\n",
    "va = torch.randn(4)\n",
    "vb = torch.randn(4)\n",
    "interp_graph({a: va, b: vb}, da=da, db=db)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8b894cc",
   "metadata": {},
   "source": [
    "In the above example, we still had variables with completely concrete\n",
    "sizes.  We can also replace all of the input sizes with dynamic sizes,\n",
    "and get a symbolic trace.  Note that in this example we gave a and\n",
    "b distinct symbolic sizes, but actually the network requires them to\n",
    "be the same and you can see the generated asserts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "2ab3be49",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "int_assert_eq(s1, s2)\n",
      "v2 = var_add(v0, v1) : (s1,)\n",
      "int_assert_eq(s1, s2)\n",
      "v3 = var_mul(v2, v1) : (s1,)\n",
      "v4 = var_sum(v3) : ()\n",
      "v5 = var_constant(val=1.0) : ()\n",
      "dv4 ------------------------\n",
      "v6 = var_expand(v5, (s1)) : (s1,)\n",
      "int_assert_eq(s1, s2)\n",
      "v7 = var_mul(v6, v1) : (s1,)\n",
      "v8 = var_mul(v6, v2) : (s1,)\n",
      "v9 = var_add(v8, v7) : (s1,)\n",
      "dv4_dv4: () = v5\n",
      "dv4_dv3: (s1,) = v6\n",
      "dv4_dv2: (s1,) = v7\n",
      "dv4_dv1: (s1,) = v9\n",
      "dv4_dv0: (s1,) = v7\n",
      "------------------------\n",
      "da = tensor([-1.0845, -1.3986,  0.4033,  0.8380])\n",
      "db = tensor([-0.6280, -3.0906, -1.3721,  2.2445])\n"
     ]
    }
   ],
   "source": [
    "reset()\n",
    "s1 = SymbolicIntNode(name=\"s1\")\n",
    "s2 = SymbolicIntNode(name=\"s2\")\n",
    "a = Variable((s1,), dtype=torch.float)\n",
    "b = Variable((s2,), dtype=torch.float)\n",
    "loss = simple(a, b).sum()\n",
    "da, db = grad(loss, [a, b])  # expand can take symbolic sizes as argument\n",
    "interp_graph({s1: 4, s2: 4, a: va, b: vb}, da=da, db=db)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41af81e0",
   "metadata": {},
   "source": [
    "In fact, with our simple interpreter, we will fail the assert EVEN if\n",
    "the original source program would have worked by broadcasting a\n",
    "one-sized dimension.  This is the sense in which true dynamic shapes\n",
    "present a different UI: dynamic shapes tell the user, \"Nuh uh, even\n",
    "if this shape is accidentally one, I'm not going to let you broadcast\n",
    "it because your code will break if you change the batch size later.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "b9d1bf15",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Failed during: int_assert_eq(s1, s2)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Traceback (most recent call last):\n",
      "  File \"/var/folders/11/bcmcs8d57q7dxbysb4w_h1ym0000gn/T/ipykernel_42390/2172691958.py\", line 2, in <cell line: 1>\n",
      "    interp_graph({s1: 1, s2: 4, a: va, b: vb}, da=da, db=db)\n",
      "  File \"/var/folders/11/bcmcs8d57q7dxbysb4w_h1ym0000gn/T/ipykernel_42390/80857361.py\", line 4, in interp_graph\n",
      "    interp_node(n, env)\n",
      "  File \"/var/folders/11/bcmcs8d57q7dxbysb4w_h1ym0000gn/T/ipykernel_42390/976668756.py\", line 27, in interp_node\n",
      "    outs = tuplify(INTERP_RULES[n.op](*args, **n.params))\n",
      "  File \"/var/folders/11/bcmcs8d57q7dxbysb4w_h1ym0000gn/T/ipykernel_42390/394099107.py\", line 13, in interp_int_assert_eq\n",
      "    assert x == y\n",
      "AssertionError\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    interp_graph({s1: 1, s2: 4, a: va, b: vb}, da=da, db=db)\n",
    "except Exception:\n",
    "    traceback.print_exc()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c9ac393",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "One of the motivating operators for dynamic shapes is nonzero (which\n",
    "is used in MaskRCNN, among other models).  It's trace implementation\n",
    "embodies the standard technique for dealing with unknown output sizes:\n",
    "allocate a fresh SymbolicIntNode and put *that* in the size of the\n",
    "returned tensor!\n",
    "\n",
    "There is a slight twist: inside the graph, it returns both a symbolic\n",
    "int (the number of nonzero elements) as well as the actual result\n",
    "tensor (containing the indices of the nonzero elements).  This is\n",
    "necessary to ensure that the symbolic int is in scope for later\n",
    "operations!  An alternative way to model this that was taken by XLA is\n",
    "to return simply the tensor, and then represent the symbolic int as a\n",
    "\"get size\" operation in the graph.  However, I like this orientation\n",
    "better, as it seems more logical (the variable depends on the symbolic\n",
    "integer, not the other way around)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "9575ea93",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "def operator_nonzero(self: Variable) -> \"Variable\":\n",
    "    s = SymbolicIntNode()\n",
    "    r = Variable((s, self.dim()), torch.long)\n",
    "    n = Node(var_nonzero_impl, (self.name,), [s.name, r.name])\n",
    "    print(f\"{n} : {r.shape}\")\n",
    "    CURRENT_GRAPH.nodes.append(n)\n",
    "    return r"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8902e9c6",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "Nonzero isn't very interesting on its own, so I also included\n",
    "implementations of advanced indexing which can make use of the\n",
    "LongTensor returned by nonzero (as well as a squeeze, as you need to\n",
    "chop off the last dimension returned by nonzero to get the output\n",
    "compatible with indexing).  In MaskRCNN, the general idea of the code\n",
    "is that you get the nonzero indexes, index them out of the tensor,\n",
    "and then compute your loss only on those entries (because they're the\n",
    "boxes that MaskRCNN actually selected for recognition!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "6f3eee16",
   "metadata": {},
   "outputs": [],
   "source": [
    "def operator_index(self: Variable, index: Variable) -> \"Variable\":\n",
    "    assert isinstance(index, Variable)  # no slices support\n",
    "    assert index.dtype == torch.long  # integer indexing only\n",
    "    assert index.dim() == 1  # 1D index only\n",
    "    r = record_var(\n",
    "        var_index, (index.shape[0],) + self.shape[1:], self.dtype, self, index\n",
    "    )\n",
    "\n",
    "    def propagate(dL_doutputs: List[Variable]):\n",
    "        (dL_dr,) = dL_doutputs\n",
    "        return [self.index_backward(index, dL_dr)]\n",
    "\n",
    "    # NB: index not recorded on tape as it is nondifferentiable\n",
    "    gradient_tape.append(\n",
    "        TapeEntry(inputs=[self.name], outputs=[r.name], propagate=propagate)\n",
    "    )\n",
    "    return r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "189f4336",
   "metadata": {},
   "outputs": [],
   "source": [
    "def operator_index_backward(\n",
    "    self: Variable, index: Variable, grad_output: Variable\n",
    ") -> \"Variable\":\n",
    "    assert isinstance(index, Variable)\n",
    "    assert index.dtype == torch.long  # integer indexing only\n",
    "    assert index.dim() == 1  # 1D index only\n",
    "    assert_int_eq(grad_output.shape[0], index.shape[0])\n",
    "    # no broadcasting\n",
    "    for i in range(1, len(self.shape)):\n",
    "        assert_int_eq(self.shape[i], grad_output.shape[i])\n",
    "    r = record_var(var_index_backward, self.shape, self.dtype, self, index, grad_output)\n",
    "\n",
    "    def propagate(dL_doutputs: List[Variable]):\n",
    "        (dL_dr,) = dL_doutputs\n",
    "        return [dL_dr[index]]\n",
    "\n",
    "    # NB: self and index not recorded as they are nondifferentiable\n",
    "    gradient_tape.append(\n",
    "        TapeEntry(inputs=[grad_output.name], outputs=[r.name], propagate=propagate)\n",
    "    )\n",
    "    return r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "d404218a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def operator_squeeze(self: Variable, dim: int) -> \"Variable\":\n",
    "    # Technically, squeeze is supposed to noop if the dimension isn't\n",
    "    # size 1.  But if the shape in question is dynamic we don't\n",
    "    # know if it is one or not.  For now, we just assert that it has to\n",
    "    # be size 1 and reduce, but technically we should use definitely_one\n",
    "    # to go between behavior\n",
    "    if not definitely_one(self.shape[dim]):\n",
    "        raise RuntimeError(\"cannot squeeze on dynamic dimension\")\n",
    "    r = record_var(\n",
    "        var_squeeze,\n",
    "        self.shape[0:dim] + self.shape[dim + 1 :],\n",
    "        self.dtype,\n",
    "        self,\n",
    "        dim=dim,\n",
    "    )\n",
    "\n",
    "    def propagate(dL_doutputs: List[Variable]):\n",
    "        (dL_dr,) = dL_outputs\n",
    "        # NB: This is only the backwards if a squeeze actually occurs\n",
    "        return [dL_dr.unsqueeze(dim)]\n",
    "\n",
    "    gradient_tape.append(\n",
    "        TapeEntry(inputs=[self.name], outputs=[r.name], propagate=propagate)\n",
    "    )\n",
    "    return r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "dd0f2be0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def operator_unsqueeze():\n",
    "    assert_int_eq(self.shape[dim], 1)\n",
    "    r = record_var(\n",
    "        var_unsqueeze,\n",
    "        self.shape[0:dim] + (1,) + self.shape[dim:],\n",
    "        self.dtype,\n",
    "        self,\n",
    "        dim=dim,\n",
    "    )\n",
    "\n",
    "    def propagate(dL_doutputs: List[Variable]):\n",
    "        (dL_dr,) = dL_outputs\n",
    "        return [dL_dr.squeeze(dim)]\n",
    "\n",
    "    gradient_tape.append(\n",
    "        TapeEntry(inputs=[self.name], outputs=[r.name], propagate=propagate)\n",
    "    )\n",
    "    return r"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ad9f0f1",
   "metadata": {},
   "source": [
    "As a warmup, let's show that indexing works."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "b2109f74",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "v2 = var_index(v0, v1) : (4, 3)\n",
      "v3 = var_sum(v2) : ()\n",
      "v4 = var_constant(val=1.0) : ()\n",
      "dv3 ------------------------\n",
      "v5 = var_expand(v4, (4, 3)) : (4, 3)\n",
      "v6 = var_index_backward(v0, v1, v5) : (2, 3)\n",
      "dv3_dv3: () = v4\n",
      "dv3_dv2: (4, 3) = v5\n",
      "dv3_dv0: (2, 3) = v6\n",
      "------------------------\n",
      "da = tensor([[3., 3., 3.],\n",
      "        [1., 1., 1.]])\n"
     ]
    }
   ],
   "source": [
    "reset()\n",
    "a = Variable((2, 3), dtype=torch.float)\n",
    "i = Variable((4,), dtype=torch.long)\n",
    "loss = a[i].sum()\n",
    "(da,) = grad(loss, [a])\n",
    "interp_graph({a: torch.randn(2, 3), b: torch.tensor([0, 0, 0, 1])}, da=da)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "585e7a42",
   "metadata": {},
   "source": [
    "Now, let's do a nontrivial symbolic case, where we index based on the\n",
    "result of nonzero"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "42a68b60",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "i0, v1 = var_nonzero_impl(v0) : (i0, 1)\n",
      "v2 = squeeze(v1, dim=1) : (i0,)\n",
      "v3 = var_index(v0, v2) : (i0,)\n",
      "L0 = var_sum(v3) : ()\n",
      "v4 = var_constant(val=1.0) : ()\n",
      "dL0 ------------------------\n",
      "v5 = var_expand(v4, (i0)) : (i0,)\n",
      "v6 = var_index_backward(v0, v2, v5) : (6,)\n",
      "dL0_dL0: () = v4\n",
      "dL0_dv3: (i0,) = v5\n",
      "dL0_dv0: (6,) = v6\n",
      "------------------------\n",
      "da = tensor([0., 1., 0., 1., 0., 0.])\n"
     ]
    }
   ],
   "source": [
    "reset()\n",
    "a = Variable((6,), dtype=torch.float)\n",
    "i = a.nonzero().squeeze(1)\n",
    "loss = a[i].sum(name=\"L0\")\n",
    "(da,) = grad(loss, [a])\n",
    "interp_graph({a: torch.clamp(torch.randn(6), min=0)}, da=da)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b709523",
   "metadata": {},
   "source": [
    "I didn't finish everything that I wanted to in this prototype.  Here\n",
    "are more things that could be done:\n",
    "\n",
    "  - We have a relatively complicated design for non-refcounted SymInt\n",
    "    in C++.  This design could be implemented here to get a better\n",
    "    understanding of how explicit reference counting affects the\n",
    "    user experience.\n",
    "\n",
    "  - Right now, the symbolic traces represent add/mul as their\n",
    "    broadcasting versions.  It is easy to tweak it using\n",
    "    `definitely_one` to explicitly represent the broadcasting\n",
    "    using an expand first.  Note that the expand may be a no-op;\n",
    "    if I broadcast (d0,) and (1,) together, I need to expand\n",
    "    (1,) to (d0,); but this is a no-op if d0 happens to be 1\n",
    "    (I'm stil obligated to record this expand though!)\n",
    "\n",
    "  - XLA's dynamic shape support also tracks upper bounds for all\n",
    "    symbolic sizes.  Incorporate support for that here.\n",
    "\n",
    "  - Incorporate some level of symbolic reasoning, improving the\n",
    "    precision of `definitely_one` or letting us check the validity\n",
    "    asserts while tracing (and not only just at runtime).  A\n",
    "    good start would be unification, but backending to Z3 would also\n",
    "    be interesting.\n",
    "\n",
    "  - Add some operators with nontrivial input/output shape\n",
    "    relationships (e.g., addition, division, etc.)"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,py:light"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}