simple_functorch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1cb2ffc9",
   "metadata": {},
   "source": [
    "This notebook walks through a self-contained implementation of\n",
    "functorch, including support for both vjp and vmap combinators (using\n",
    "PyTorch only to implement primitive tensor operations).  It follows\n",
    "the tradition of\n",
    "[Autodidax](https://jax.readthedocs.io/en/latest/autodidax.html) (a\n",
    "pedagogical reimplementation of JAX, the library functorch is inspired\n",
    "by) and [Simple\n",
    "Autograd](https://colab.research.google.com/drive/1VpeE6UvEPRz9HmsHh1KS0XxXjYu533EC?usp=sharing)\n",
    "(Zachary Devito's pedagogical reimplementation of autograd, which the\n",
    "autograd system in this notebook is based off of.) You can [open this\n",
    "file in\n",
    "Colab](https://colab.research.google.com/github/albanD/subclass_zoo/blob/main/simple_functorch.ipynb)\n",
    "and play around with the examples.\n",
    "\n",
    "As a simplified implementation of functorch, this notebook also makes\n",
    "it easier to investigate some more subtle aspects of how PyTorch's\n",
    "native autograd system interacts with composable transforms.  In\n",
    "particular, we will see that PyTorch's native implementation of double\n",
    "backwards (which shares the same tape through multiple levels of\n",
    "differentiation) differs from functorch's nested grad implementation\n",
    "(which maintains a separate tape per level)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ed150b0",
   "metadata": {},
   "source": [
    "To get started, we replicate some of the data structures and helper functions\n",
    "from Simple Autograd."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 236,
   "id": "ba464220",
   "metadata": {
    "lines_to_end_of_cell_marker": 2
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import Tensor\n",
    "from typing import List, NamedTuple, Callable, Dict, Optional\n",
    "import contextlib\n",
    "import functools\n",
    "from dataclasses import dataclass\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[Tensor]], List[Tensor]]\n",
    "\n",
    "\n",
    "_name = 0\n",
    "\n",
    "\n",
    "def fresh_name() -> str:\n",
    "    \"\"\"create a new unique name for a variable: v0, v1, v2\"\"\"\n",
    "    global _name\n",
    "    r = f\"v{_name}\"\n",
    "    _name += 1\n",
    "    return r"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0af3ca0b",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "This is a little helper function for converting the dim argument in\n",
    "sum into an explicit list of dimensions that will be reduced over.\n",
    "It takes the dim of the tensor we are summing over and the dim\n",
    "argument itself."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 237,
   "id": "836a1635",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "def sum_dims(*, input_dim, dim):\n",
    "    if dim is None:\n",
    "        return tuple(range(0, input_dim))\n",
    "    elif isinstance(dim, int):\n",
    "        return (dim,)\n",
    "    else:\n",
    "        return tuple(sorted(dim))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f74fb48",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "In Simple Autograd, we provided a Variable wrapper class which\n",
    "provided a traditional Tensor style interface for our objects; in\n",
    "functorch proper, objects are repeatedly wrapped in this way to\n",
    "implement multipler layers of transformations.\n",
    "\n",
    "In my opinion, this sort of wrapper makes it more difficult to\n",
    "understand the flow of logic.  So in Simple Functorch, we take a\n",
    "different approach: we won't make use of a wrapper class at all,\n",
    "instead showing how to add it in the end as syntax sugar on top of our\n",
    "system.\n",
    "\n",
    "For debuggability purposes, however, it is nice to have a way to\n",
    "identify variables by a human readable name.  We'll do this by setting\n",
    "a t_name attribute on PyTorch tensors whenever we allocate a new\n",
    "tensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 238,
   "id": "f9d17aef",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "def label(t: Tensor, name: str = None):\n",
    "    if not hasattr(t, \"t_name\"):\n",
    "        t.t_name = name or fresh_name()\n",
    "    return t"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd046576",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "So if we aren't going to have a wrapper around each tensor, how will\n",
    "we actually implement our logic?  We will organize our various layers\n",
    "of transformations as separate Dispatcher objects, which define\n",
    "methods for performing operations on tensors, but are not Tensors\n",
    "themselves.  For example, instead of defining Tensor.add(Tensor), we\n",
    "will define Dispatcher.add(Tensor, Tensor).  If you are familiar with\n",
    "historical ATen, in the original implementation of ATen, these\n",
    "correspond to the CPUType/CUDAType/VariableType objects from that\n",
    "implementation (this was replaced with the modern dispatcher as the\n",
    "original vtable-based implementation did not support adding custom\n",
    "operators.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 239,
   "id": "59a69d9d",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "class Dispatcher:\n",
    "    inner = None\n",
    "\n",
    "    def mul(self, lhs, rhs):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def add(self, lhs, rhs):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    # Sum has been generalized to take an optional dim argument, which we\n",
    "    # will need for Batched tensors\n",
    "    def sum(self, input, dim=None, name=None):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def expand(self, input, sizes):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    # For closure under Batched tensors, we need these operations...\n",
    "    def unsqueeze(self, input, dim):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def squeeze(self, input, dim):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    # ...and we also need to overload the meaning of size/ones to\n",
    "    # hide/reinsert batch dimensions.  We also introduce a concept\n",
    "    # of \"lifting\" a tensor to be batched by broadcasting it on\n",
    "    # a dimension\n",
    "    def size(self, input):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def ones(self, size):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def lift(self, input, d):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    # For convenience, we provide dim, which just returns the length of\n",
    "    # the sizes\n",
    "    def dim(self, input):\n",
    "        return len(self.size(input))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0e21984",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "To start with, we can implement a backend dispatcher layer Torch,\n",
    "which just forwards the operator calls to our underlying library PyTorch\n",
    "(and ensures that all the allocated tensors are labeled with variable).  You could\n",
    "also imagine replacing this with a Numpy backend or even a pure Python\n",
    "variant (although this file is not currently setup to do so.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 240,
   "id": "27101073",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "class Torch(Dispatcher):\n",
    "    def mul(self, lhs, rhs):\n",
    "        return label(torch.mul(lhs, rhs))\n",
    "\n",
    "    def add(self, lhs, rhs):\n",
    "        return label(torch.add(lhs, rhs))\n",
    "\n",
    "    def sum(self, input, dim=None, name=None):\n",
    "        if dim is None:\n",
    "            return label(torch.sum(input), name)\n",
    "        else:\n",
    "            return label(torch.sum(input, dim), name)\n",
    "\n",
    "    def expand(self, input, sizes):\n",
    "        return label(input.expand(sizes))\n",
    "\n",
    "    def unsqueeze(self, input, dim):\n",
    "        return label(torch.unsqueeze(input, dim))\n",
    "\n",
    "    def squeeze(self, input, dim):\n",
    "        return label(torch.squeeze(input, dim))\n",
    "\n",
    "    def size(self, input):\n",
    "        # Return size a tuple for marginally more compact printing\n",
    "        assert isinstance(input, torch.Tensor)\n",
    "        return tuple(input.size())\n",
    "\n",
    "    def ones(self, size):\n",
    "        return label(torch.ones(size))\n",
    "\n",
    "    def lift(self, input, d):\n",
    "        assert d is self\n",
    "        return input"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e428f835",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "Dispatcher layers are composable via object composition: we can\n",
    "imagine a stack of dispatchers, each one calling into the next.\n",
    "For example, the Logger dispatcher simply prints out what operation\n",
    "was called on it, and then forwards on the operation to the inner\n",
    "dispatcher."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 241,
   "id": "1500269c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Logger(Dispatcher):\n",
    "    def __init__(self, inner, *, name):\n",
    "        self.inner = inner\n",
    "        self.name = f\"  {name}\"\n",
    "\n",
    "    def size(self, input):\n",
    "        # don't log size calls\n",
    "        return self.inner.size(input)\n",
    "\n",
    "    def ones(self, size):\n",
    "        r = self.inner.ones(size)\n",
    "        print(f\"{self.name} {r.t_name}: {self.size(r)} = ones({size})\")\n",
    "        return r\n",
    "\n",
    "    def mul(self, lhs, rhs):\n",
    "        r = self.inner.mul(lhs, rhs)\n",
    "        if isinstance(rhs, float):\n",
    "            print(f\"{self.name} {r.t_name}: {self.size(r)} = {lhs.t_name} * {rhs}\")\n",
    "        else:\n",
    "            print(\n",
    "                f\"{self.name} {r.t_name}: {self.size(r)} = {lhs.t_name} * {rhs.t_name}\"\n",
    "            )\n",
    "        return r\n",
    "\n",
    "    def add(self, lhs, rhs):\n",
    "        r = self.inner.add(lhs, rhs)\n",
    "        print(f\"{self.name} {r.t_name}: {self.size(r)} = {lhs.t_name} + {rhs.t_name}\")\n",
    "        return r\n",
    "\n",
    "    def sum(self, input, dim=None, name=None):\n",
    "        r = self.inner.sum(input, dim=dim, name=name)\n",
    "        print(f\"{self.name} {r.t_name}: {self.size(r)} = {input.t_name}.sum(dim={dim})\")\n",
    "        return r\n",
    "\n",
    "    def unsqueeze(self, input, dim):\n",
    "        r = self.inner.unsqueeze(input, dim)\n",
    "        print(\n",
    "            f\"{self.name} {r.t_name}: {self.size(r)} = {input.t_name}.unsqueeze({dim})\"\n",
    "        )\n",
    "        return r\n",
    "\n",
    "    def squeeze(self, input, dim):\n",
    "        r = self.inner.squeeze(input, dim)\n",
    "        print(f\"{self.name} {r.t_name}: {self.size(r)} = {input.t_name}.squeeze({dim})\")\n",
    "        return r\n",
    "\n",
    "    def expand(self, input, sizes):\n",
    "        r = self.inner.expand(input, sizes)\n",
    "        print(\n",
    "            f\"{self.name} {r.t_name}: {self.size(r)} = {input.t_name}.expand({sizes})\"\n",
    "        )\n",
    "        return r\n",
    "\n",
    "    def lift(self, input, d):\n",
    "        if d is self:\n",
    "            return input\n",
    "        return self.inner.lift(input, d)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e323b552",
   "metadata": {},
   "source": [
    "Here is a simple example of using Logger and Torch together.  Whenever\n",
    "we make calls to operations, we must do so via the Dispatcher object.\n",
    "We will explicitly write out all of these calls before we add wrapper\n",
    "class sugaring."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 242,
   "id": "26b4c5e5",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Torch v0: (2,) = ones(2)\n",
      "  Torch v1: (2,) = ones(2)\n",
      "  Torch v2: (2,) = v0 + v1\n",
      "tensor([2., 2.])\n"
     ]
    }
   ],
   "source": [
    "d = Logger(Torch(), name=\"Torch\")\n",
    "print(d.add(d.ones(2), d.ones(2)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1dcd422b",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "With the Dispatcher structure in hand, we are now in a good place to\n",
    "port the autograd implementation from Simple Autograd into our new\n",
    "framework."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 243,
   "id": "54293b9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Autograd(Dispatcher):\n",
    "    # create_graph here corresponds to the create_graph kwarg in traditional\n",
    "    # PyTorch, which controls whether or not the graph of the derivative\n",
    "    # will be constructed, allowing computing higher order derivatives.\n",
    "    # We will see that although create_graph=True allows Autograd to directly\n",
    "    # support higher order derivatives, layering an Autograd to another\n",
    "    # Autograd will also allow higher order derivatives.\n",
    "    def __init__(self, inner, *, name=\"Autograd\", create_graph: bool = False):\n",
    "        self.inner = inner\n",
    "        self.gradient_tape = []\n",
    "        self.name = name\n",
    "        self.create_graph = create_graph\n",
    "\n",
    "    # create_graph controls where add/mul/etc calls from the backwards\n",
    "    # propagators go: if you create_graph, they recursively call back to\n",
    "    # the current Autograd dispatcher; otherwise they move on to the inner\n",
    "    # layer.\n",
    "    def backward_inner(self):\n",
    "        if self.create_graph:\n",
    "            return self\n",
    "        else:\n",
    "            return self.inner\n",
    "\n",
    "    def mul(self, lhs, rhs):\n",
    "        if isinstance(rhs, float) and rhs == 1.0:\n",
    "            # peephole optimization\n",
    "            return lhs\n",
    "\n",
    "        # define forward\n",
    "        # first, run the operation in the inner layer to get the initial\n",
    "        # result\n",
    "        r = self.inner.mul(lhs, rhs)\n",
    "        # We directly implement printing here as it indicates whether or not\n",
    "        # this operation was saved to the tape or not\n",
    "        print(f\"{self.name} {r.t_name}: {self.size(r)} = {lhs.t_name} * {rhs.t_name}\")\n",
    "\n",
    "        # record what the inputs and outputs of the op were\n",
    "        inputs = [lhs.t_name, rhs.t_name]\n",
    "        outputs = [r.t_name]\n",
    "\n",
    "        # define backprop\n",
    "        def propagate(dL_doutputs: List[Tensor]):\n",
    "            (dL_dr,) = dL_doutputs\n",
    "\n",
    "            dr_dlhs = rhs  # partial derivative of r = lhs*rhs\n",
    "            dr_drhs = lhs  # partial derivative of r = lhs*rhs\n",
    "\n",
    "            # chain rule propagation from outputs to inputs of multiply.\n",
    "            # Notice that the propagation rule may itself call\n",
    "            # other operations; depending on create_graph, they may\n",
    "            # either go to self or self.inner; self.backward_inner()\n",
    "            # controls which one we go to.\n",
    "            dL_dlhs = self.backward_inner().mul(dL_dr, dr_dlhs)\n",
    "            dL_drhs = self.backward_inner().mul(dL_dr, dr_drhs)\n",
    "            dL_dinputs = [dL_dlhs, dL_drhs]\n",
    "            return dL_dinputs\n",
    "\n",
    "        # finally, we record the compute we did on the tape\n",
    "        self.gradient_tape.append(\n",
    "            TapeEntry(inputs=inputs, outputs=outputs, propagate=propagate)\n",
    "        )\n",
    "        return r\n",
    "\n",
    "    # The rest of the implementations follow in the same way and can\n",
    "    # be skipped\n",
    "\n",
    "    def add(self, lhs, rhs):\n",
    "        # Add follows a similar pattern to Mul, but it doesn't end up\n",
    "        # capturing any variables.\n",
    "        r = self.inner.add(lhs, rhs)\n",
    "        print(f\"{self.name} {r.t_name}: {self.size(r)} = {lhs.t_name} + {rhs.t_name}\")\n",
    "\n",
    "        def propagate(dL_doutputs: List[Tensor]):\n",
    "            (dL_dr,) = dL_doutputs\n",
    "            dr_dlhs = 1.0\n",
    "            dr_drhs = 1.0\n",
    "            dL_dlhs = self.backward_inner().mul(dL_dr, dr_dlhs)\n",
    "            dL_drhs = self.backward_inner().mul(dL_dr, dr_drhs)\n",
    "            return [dL_dlhs, dL_drhs]\n",
    "\n",
    "        self.gradient_tape.append(\n",
    "            TapeEntry(\n",
    "                inputs=[lhs.t_name, rhs.t_name], outputs=[r.t_name], propagate=propagate\n",
    "            )\n",
    "        )\n",
    "        return r\n",
    "\n",
    "    # Extended to handle dim argument for Batched (later)\n",
    "    def sum(self, input: Tensor, dim=None, name: Optional[str] = None):\n",
    "        r = self.inner.sum(input, dim=dim, name=name)\n",
    "        print(f\"{self.name} {r.t_name}: {self.size(r)} = {input.t_name}.sum(dim={dim})\")\n",
    "\n",
    "        def propagate(dL_doutputs: List[Tensor]):\n",
    "            (dL_dr,) = dL_doutputs\n",
    "            size = self.inner.size(input)\n",
    "            res = dL_dr\n",
    "            # Broadcast over all dimensions that were reduced over\n",
    "            for i in sum_dims(input_dim=self.inner.dim(input), dim=dim):\n",
    "                res = self.backward_inner().unsqueeze(res, i)\n",
    "            return [self.backward_inner().expand(res, size)]\n",
    "\n",
    "        self.gradient_tape.append(\n",
    "            TapeEntry(inputs=[input.t_name], outputs=[r.t_name], propagate=propagate)\n",
    "        )\n",
    "        return r\n",
    "\n",
    "    # Unlike Simple Autograd, this expand requires the input to have\n",
    "    # been unsqueezed before hand.  This lets us avoid having to do\n",
    "    # at::sum_to for the nontrivial case (which is more complicated)\n",
    "    def expand(self, input: Tensor, sizes: List[int]):\n",
    "        assert self.inner.dim(input) == len(sizes)  # only works if dims match\n",
    "        r = self.inner.expand(input, sizes)\n",
    "        print(\n",
    "            f\"{self.name} {r.t_name}: {self.size(r)} = {input.t_name}.expand({sizes})\"\n",
    "        )\n",
    "\n",
    "        def propagate(dL_doutputs: List[Tensor]):\n",
    "            (dL_dr,) = dL_doutputs\n",
    "            input_size = self.inner.size(input)\n",
    "            dims = tuple(\n",
    "                i for i in range(self.inner.dim(input)) if input_size[i] != sizes[i]\n",
    "            )\n",
    "            # We wanted a sum keepdim=True, but I didn't want to force\n",
    "            # everyone to support it so manually unsqueeze\n",
    "            res = self.backward_inner().sum(dL_dr, dims)\n",
    "            for d in dims:\n",
    "                res = self.backward_inner().unsqueeze(res, d)\n",
    "            return [res]\n",
    "\n",
    "        self.gradient_tape.append(\n",
    "            TapeEntry(inputs=[input.t_name], outputs=[r.t_name], propagate=propagate)\n",
    "        )\n",
    "        return r\n",
    "\n",
    "    # Unsqueeze are required for sum backwards, and then squeeze is required\n",
    "    # for closure.  Size needed for batched tensor to modify size.\n",
    "\n",
    "    def size(self, input: Tensor):\n",
    "        return self.inner.size(input)\n",
    "\n",
    "    def squeeze(self, input: Tensor, dim):\n",
    "        r = self.inner.squeeze(input, dim)\n",
    "        print(\n",
    "            f\"{self.name} {r.t_name}: {self.size(r)} = {input.t_name}.squeeze(dim={dim})\"\n",
    "        )\n",
    "\n",
    "        def propagate(dL_doutputs: List[Tensor]):\n",
    "            (dL_dr,) = dL_outputs\n",
    "            return [self.backward_inner().unsqueeze(dL_dr, dim)]\n",
    "\n",
    "        self.gradient_tape.append(\n",
    "            TapeEntry(inputs=[input.t_name], outputs=[r.t_name], propagate=propagate)\n",
    "        )\n",
    "        return r\n",
    "\n",
    "    def unsqueeze(self, input: Tensor, dim):\n",
    "        r = self.inner.unsqueeze(input, dim)\n",
    "        print(\n",
    "            f\"{self.name} {r.t_name}: {self.size(r)} = {input.t_name}.unsqueeze(dim={dim})\"\n",
    "        )\n",
    "\n",
    "        def propagate(dL_doutputs: List[Tensor]):\n",
    "            (dL_dr,) = dL_doutputs\n",
    "            return [self.backward_inner().squeeze(dL_dr, dim)]\n",
    "\n",
    "        self.gradient_tape.append(\n",
    "            TapeEntry(inputs=[input.t_name], outputs=[r.t_name], propagate=propagate)\n",
    "        )\n",
    "        return r\n",
    "\n",
    "    def ones(self, size):\n",
    "        return self.inner.ones(size)\n",
    "\n",
    "    def lift(self, input, d):\n",
    "        if d is self:\n",
    "            return input\n",
    "        return self.inner.lift(input, d)\n",
    "\n",
    "    def grad(self, L, desired_results: List[Tensor]) -> List[Tensor]:\n",
    "        # this map holds dL/dX for all values X\n",
    "        dL_d: Dict[str, Tensor] = {}\n",
    "        # It starts by initializing the 'seed' dL/dL, which is 1\n",
    "        # TODO: indirect this via the backend\n",
    "        dL_d[L.t_name] = self.inner.ones(())\n",
    "        print(f\"-- {self.name} d{L.t_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(self.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] = self.backward_inner().add(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'{self.name} d{L.t_name}_d{name} = {value.t_name}')\n",
    "        print(f\"------------------------\")\n",
    "\n",
    "        return gather_grad(desired.t_name for desired in desired_results)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "901c6b4a",
   "metadata": {},
   "source": [
    "To calculate some simple gradients, we can compose Autograd with\n",
    "Torch and get the result we expect."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 244,
   "id": "b725c318",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Autograd v5: (4,) = v3 + v4\n",
      "Autograd v6: (4,) = v5 * v4\n",
      "-- Autograd dv6 -------\n",
      "------------------------\n",
      "da tensor([0.3074, 0.6341, 0.4901, 0.8964])\n",
      "db tensor([1.1111, 2.0364, 1.0687, 1.9249])\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "a, b = label(torch.rand(4)), label(torch.rand(4))\n",
    "\n",
    "\n",
    "def simple(d, a, b):\n",
    "    t = d.add(a, b)\n",
    "    return d.mul(t, b)\n",
    "\n",
    "\n",
    "d = Autograd(Torch())\n",
    "\n",
    "loss = simple(d, a, b)\n",
    "da, db = d.grad(loss, [a, b])\n",
    "print(\"da\", da)\n",
    "print(\"db\", db)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fdd64990",
   "metadata": {},
   "source": [
    "To compute higher order gradients, we have two options.  First,\n",
    "we can do traditional PyTorch style higher order differentiation\n",
    "with `create_graph=True`, writing the backpropagation computations directly\n",
    "into the tape so they can be further differentiated over.  This is also\n",
    "what the original Simple Autograd implementation does."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 245,
   "id": "7e1a5342",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Autograd v13: (4,) = v3 + v4\n",
      "Autograd v14: (4,) = v13 * v4\n",
      "Autograd L0: () = v14.sum(dim=None)\n",
      "-- Autograd dL0 -------\n",
      "Autograd v16: (1,) = v15.unsqueeze(dim=0)\n",
      "Autograd v17: (4,) = v16.expand((4,))\n",
      "Autograd v18: (4,) = v17 * v4\n",
      "Autograd v19: (4,) = v17 * v13\n",
      "Autograd v20: (4,) = v19 + v18\n",
      "------------------------\n",
      "Autograd v21: (4,) = v18 * v18\n",
      "Autograd v22: (4,) = v20 * v20\n",
      "Autograd v23: (4,) = v21 + v22\n",
      "Autograd L1: () = v23.sum(dim=None)\n",
      "-- Autograd dL1 -------\n",
      "Autograd v25: (1,) = v24.unsqueeze(dim=0)\n",
      "Autograd v26: (4,) = v25.expand((4,))\n",
      "Autograd v27: (4,) = v26 * v20\n",
      "Autograd v28: (4,) = v26 * v20\n",
      "Autograd v29: (4,) = v27 + v28\n",
      "Autograd v30: (4,) = v26 * v18\n",
      "Autograd v31: (4,) = v26 * v18\n",
      "Autograd v32: (4,) = v30 + v31\n",
      "Autograd v33: (4,) = v32 + v29\n",
      "Autograd v34: (4,) = v29 * v13\n",
      "Autograd v35: (4,) = v29 * v17\n",
      "Autograd v36: (4,) = v33 * v4\n",
      "Autograd v37: (4,) = v33 * v17\n",
      "Autograd v38: (4,) = v34 + v36\n",
      "Autograd v39: () = v38.sum(dim=(0,))\n",
      "Autograd v40: (1,) = v39.unsqueeze(dim=0)\n",
      "Autograd v41: () = v40.squeeze(dim=0)\n",
      "Autograd v42: (4,) = v37 + v35\n",
      "------------------------\n",
      "da tensor([2.2222, 4.0728, 2.1373, 3.8498])\n",
      "db tensor([5.0593, 9.4137, 5.2548, 9.4926])\n"
     ]
    }
   ],
   "source": [
    "d = Autograd(Torch(), create_graph=True)\n",
    "\n",
    "# I slightly generalized this function so that it works for the next\n",
    "# example; d2 is the dispatcher run on the first grad call, and d1 is\n",
    "# for the second (we'll see why the numbers are inverted shortly).\n",
    "def run_gradients(d2, d1):\n",
    "    # our first loss\n",
    "    L0 = d2.sum(simple(d2, a, b), name=\"L0\")\n",
    "\n",
    "    # compute derivatives of our inputs\n",
    "    dL0_da, dL0_db = d2.grad(L0, [a, b])\n",
    "\n",
    "    # In real code, how would we switch from executing from d2 to d1?\n",
    "    # In functorch, the d2 dispatch calls would happen in the inside of\n",
    "    # a higher-order grad() call; when we exit from this call, all\n",
    "    # of the involved tensors are unwrapped.\n",
    "\n",
    "    # now lets compute the L2 norm of our derivatives\n",
    "    L1 = d1.sum(d1.add(d1.mul(dL0_da, dL0_da), d1.mul(dL0_db, dL0_db)), name=\"L1\")\n",
    "\n",
    "    # and take the gradient of that.\n",
    "    # notice there are two losses involved1.\n",
    "    dL1_da, dL1_db = d1.grad(L1, [a, b])\n",
    "    return dL1_da, dL1_db\n",
    "\n",
    "\n",
    "da, db = run_gradients(d, d)\n",
    "print(\"da\", da)\n",
    "print(\"db\", db)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b17d02c",
   "metadata": {},
   "source": [
    "Our second option is to follow functorch's implementation strategy, which\n",
    "is to stack two Autograd dispatchers on top of each other.  Here, it is\n",
    "not necessary to `create_graph=True`, because when the backpropagator forwards\n",
    "to the inner dispatcher, it will record those operations on the tape too.\n",
    "But if you look at the output, you will notice something very interesting:\n",
    "the first portion of the tape is exactly replicated between Autograd1 and\n",
    "Autograd2: we're duplicating the tape in this case!  So PyTorch's default\n",
    "implementation of backwards is more efficient, because it avoids having to\n",
    "record the tape twice (although this doesn't matter too much, because the\n",
    "saved tensors themselves can be shared between the two tapes, so it is just\n",
    "the operator graph that is duplicated.\n",
    "\n",
    "This is our first example of using two dispatchers.  While we are\n",
    "performing the inner grad, we perform our operations on the outer\n",
    "dispatcher `d2`; after we are done with the inner grad we switch to\n",
    "`d1`.  Intuitively, this corresponds from passing out of the inner\n",
    "`grad` call to the outer `grad` call."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 246,
   "id": "a7c47f5e",
   "metadata": {
    "lines_to_end_of_cell_marker": 0,
    "lines_to_next_cell": 1
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Torch v43: (4,) = v3 + v4\n",
      "Autograd1 v43: (4,) = v3 + v4\n",
      "Autograd2 v43: (4,) = v3 + v4\n",
      "  Torch v44: (4,) = v43 * v4\n",
      "Autograd1 v44: (4,) = v43 * v4\n",
      "Autograd2 v44: (4,) = v43 * v4\n",
      "  Torch L0: () = v44.sum(dim=None)\n",
      "Autograd1 L0: () = v44.sum(dim=None)\n",
      "Autograd2 L0: () = v44.sum(dim=None)\n",
      "  Torch v45: () = ones(())\n",
      "-- Autograd2 dL0 -------\n",
      "  Torch v46: (1,) = v45.unsqueeze(0)\n",
      "Autograd1 v46: (1,) = v45.unsqueeze(dim=0)\n",
      "  Torch v47: (4,) = v46.expand((4,))\n",
      "Autograd1 v47: (4,) = v46.expand((4,))\n",
      "  Torch v48: (4,) = v47 * v4\n",
      "Autograd1 v48: (4,) = v47 * v4\n",
      "  Torch v49: (4,) = v47 * v43\n",
      "Autograd1 v49: (4,) = v47 * v43\n",
      "  Torch v50: (4,) = v49 + v48\n",
      "Autograd1 v50: (4,) = v49 + v48\n",
      "------------------------\n",
      "  Torch v51: (4,) = v48 * v48\n",
      "Autograd1 v51: (4,) = v48 * v48\n",
      "  Torch v52: (4,) = v50 * v50\n",
      "Autograd1 v52: (4,) = v50 * v50\n",
      "  Torch v53: (4,) = v51 + v52\n",
      "Autograd1 v53: (4,) = v51 + v52\n",
      "  Torch L1: () = v53.sum(dim=None)\n",
      "Autograd1 L1: () = v53.sum(dim=None)\n",
      "  Torch v54: () = ones(())\n",
      "-- Autograd1 dL1 -------\n",
      "  Torch v55: (1,) = v54.unsqueeze(0)\n",
      "  Torch v56: (4,) = v55.expand((4,))\n",
      "  Torch v57: (4,) = v56 * 1.0\n",
      "  Torch v58: (4,) = v56 * 1.0\n",
      "  Torch v59: (4,) = v58 * v50\n",
      "  Torch v60: (4,) = v58 * v50\n",
      "  Torch v61: (4,) = v59 + v60\n",
      "  Torch v62: (4,) = v57 * v48\n",
      "  Torch v63: (4,) = v57 * v48\n",
      "  Torch v64: (4,) = v62 + v63\n",
      "  Torch v65: (4,) = v61 * 1.0\n",
      "  Torch v66: (4,) = v61 * 1.0\n",
      "  Torch v67: (4,) = v64 + v66\n",
      "  Torch v68: (4,) = v65 * v43\n",
      "  Torch v69: (4,) = v65 * v47\n",
      "  Torch v70: (4,) = v67 * v4\n",
      "  Torch v71: (4,) = v67 * v47\n",
      "  Torch v72: (4,) = v68 + v70\n",
      "  Torch v73: () = v72.sum(dim=(0,))\n",
      "  Torch v74: (1,) = v73.unsqueeze(0)\n",
      "  Torch v75: () = v74.squeeze(0)\n",
      "  Torch v76: (4,) = v69 * 1.0\n",
      "  Torch v77: (4,) = v69 * 1.0\n",
      "  Torch v78: (4,) = v71 + v77\n",
      "------------------------\n",
      "da tensor([2.2222, 4.0728, 2.1373, 3.8498])\n",
      "db tensor([5.0593, 9.4137, 5.2548, 9.4926])\n"
     ]
    }
   ],
   "source": [
    "# turning off create_graph will impede us from seeing the logging lines for\n",
    "# the second backwards, so we turn on logging for Torch to see them\n",
    "d1 = Autograd(Logger(Torch(), name=\"Torch\"), name=\"Autograd1\", create_graph=False)\n",
    "d2 = Autograd(d1, name=\"Autograd2\", create_graph=False)\n",
    "\n",
    "da, db = run_gradients(d2, d1)\n",
    "print(\"da\", da)\n",
    "print(\"db\", db)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f343af81",
   "metadata": {},
   "source": [
    "Under what situations might it be profitable to keep the two tapes separate?\n",
    "One guess we might have is if there is another functional transformation\n",
    "wedged between the two autograd transformations.  We would then expect the\n",
    "backwards formula we save to be different between the two tapes.  To do this, I\n",
    "first need to implement batched tensors.\n",
    "\n",
    "One unusual thing about this implementation is that we do not need to wrap\n",
    "tensors to change their sizes; instead, we just override the meaning of\n",
    "size() on the dispatcher to hide batch dimensions.  One case we do not\n",
    "exercise in this example is implicit broadcasting when you combine a tensor\n",
    "that is not batched with a tensor that is batched: without wrappers, a user\n",
    "must explicitly lift (e.g., unsqueeze and expand) tensors they wish to\n",
    "replicate across the batch dimension.  The code below will blindly attempt to\n",
    "reinterpret a tensor as a batched tensor, even when it may not make sense (if\n",
    "there is a size mismatch, however, you will get an assert failure).  Similarly,\n",
    "once you exit a vmap region, all previously vmap'ed tensors \"magically\" become\n",
    "unbatched.  functorch did not pursue this implementation because at the time\n",
    "Tensor.size() was not virtual and thus it was not possible to override (this\n",
    "will be changing soon)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 247,
   "id": "dd1d0688",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This implementation of Batched only supports inserting a dimension\n",
    "# at the very front\n",
    "class Batched(Dispatcher):\n",
    "    def __init__(self, inner, *, length, name=\"Batched\"):\n",
    "        self.inner = inner\n",
    "        self.name = name\n",
    "        self.length = length\n",
    "\n",
    "    def size(self, input):\n",
    "        sizes = self.inner.size(input)\n",
    "        assert sizes[0] == self.length\n",
    "        return sizes[1:]\n",
    "\n",
    "    def ones(self, size):\n",
    "        return self.inner.ones((self.length,) + size)\n",
    "\n",
    "    def mul(self, lhs, rhs):\n",
    "        assert self.inner.size(lhs)[0] == self.length\n",
    "        if not isinstance(rhs, float):\n",
    "            assert self.inner.size(rhs)[0] == self.length\n",
    "        return self.inner.mul(lhs, rhs)\n",
    "\n",
    "    def add(self, lhs, rhs):\n",
    "        assert self.inner.size(lhs)[0] == self.length\n",
    "        assert self.inner.size(rhs)[0] == self.length\n",
    "        return self.inner.add(lhs, rhs)\n",
    "\n",
    "    def sum(self, input, dim=None, name=None):\n",
    "        # offset all the summed over dimensions by one\n",
    "        assert self.inner.size(input)[0] == self.length\n",
    "        dim = tuple(\n",
    "            i + 1 for i in sum_dims(input_dim=self.inner.dim(input) - 1, dim=dim)\n",
    "        )\n",
    "        return self.inner.sum(input, dim, name=name)\n",
    "\n",
    "    def expand(self, input, sizes):\n",
    "        # offset sizes by one\n",
    "        assert self.inner.size(input)[0] == self.length\n",
    "        return self.inner.expand(input, (self.inner.size(input)[0],) + sizes)\n",
    "\n",
    "    def squeeze(self, input, dim):\n",
    "        # offset dim by one\n",
    "        assert self.inner.size(input)[0] == self.length\n",
    "        return self.inner.squeeze(input, dim + 1)\n",
    "\n",
    "    def unsqueeze(self, input, dim):\n",
    "        # offset dim by one\n",
    "        assert self.inner.size(input)[0] == self.length\n",
    "        return self.inner.unsqueeze(input, dim + 1)\n",
    "\n",
    "    # The lift operation takes a tensor associated with some inner\n",
    "    # dispatcher, and \"lifts\" it so that it is interpreted neutrally\n",
    "    # for the outer dispatcher.  For most dispatchers this is trivial,\n",
    "    # but for batched tensor it is not: given a tensor x, to interpret\n",
    "    # it as x under the Batching dispatcher, we have to expand it so\n",
    "    # that it is broadcasted along its first dimension.\n",
    "    def lift(self, input, d):\n",
    "        if d is self:\n",
    "            return input\n",
    "        b_input = self.inner.unsqueeze(input, 0)\n",
    "        b_input = self.inner.expand(b_input, (self.length,) + self.inner.size(input))\n",
    "        return self.inner.lift(b_input, d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 248,
   "id": "dffb9a7e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Autograd1 v81: (2, 4) = v79 + v80\n",
      "Autograd3 v81: (4,) = v79 + v80\n",
      "Autograd1 v82: (2, 4) = v81 * v80\n",
      "Autograd3 v82: (4,) = v81 * v80\n",
      "Autograd1 L0: (2,) = v82.sum(dim=(1,))\n",
      "Autograd3 L0: () = v82.sum(dim=0)\n",
      "-- Autograd3 dL0 -------\n",
      "Autograd1 v84: (2, 1) = v83.unsqueeze(dim=1)\n",
      "Autograd1 v85: (2, 4) = v84.expand((2, 4))\n",
      "Autograd1 v86: (2, 4) = v85 * v80\n",
      "Autograd1 v87: (2, 4) = v85 * v81\n",
      "Autograd1 v88: (2, 4) = v87 + v86\n",
      "------------------------\n",
      "Autograd1 v89: (2, 4) = v86 * v86\n",
      "Autograd1 v90: (2, 4) = v88 * v88\n",
      "Autograd1 v91: (2, 4) = v89 + v90\n",
      "Autograd1 L1: () = v91.sum(dim=None)\n",
      "-- Autograd1 dL1 -------\n",
      "------------------------\n",
      "va tensor([[0.4556, 0.6323, 0.3489, 0.4017],\n",
      "        [0.0223, 0.1689, 0.2939, 0.5185]])\n",
      "vb tensor([[0.6977, 0.8000, 0.1610, 0.2823],\n",
      "        [0.6816, 0.9152, 0.3971, 0.8742]])\n",
      "dva tensor([[3.7019, 4.4647, 1.3419, 1.9325],\n",
      "        [2.7711, 3.9985, 2.1762, 4.5337]])\n",
      "dvb tensor([[ 8.7992, 10.5293,  3.0059,  4.4296],\n",
      "        [ 6.9054,  9.8274,  5.1466, 10.8156]])\n"
     ]
    }
   ],
   "source": [
    "# Our inputs are batched this time!\n",
    "va, vb = label(torch.rand(2, 4)), label(torch.rand(2, 4))\n",
    "\n",
    "d1 = Autograd(Torch(), name=\"Autograd1\", create_graph=False)\n",
    "d2 = Batched(d1, length=2, name=\"Batched2\")\n",
    "d3 = Autograd(d2, name=\"Autograd3\", create_graph=False)\n",
    "\n",
    "\n",
    "def run_batched_gradients(d3, d2, d1):\n",
    "    # our first loss\n",
    "    # we write the dimension we reduce on explicitly for clarity\n",
    "    L0 = d3.sum(simple(d3, va, vb), dim=0, name=\"L0\")\n",
    "\n",
    "    # compute derivatives of our inputs\n",
    "    dL0_da, dL0_db = d3.grad(L0, [va, vb])\n",
    "\n",
    "    # now lets compute the L2 norm of our derivatives\n",
    "    L1 = d1.sum(d1.add(d1.mul(dL0_da, dL0_da), d1.mul(dL0_db, dL0_db)), name=\"L1\")\n",
    "\n",
    "    # and take the gradient of that.\n",
    "    # notice there are two losses involved1.\n",
    "    dL1_da, dL1_db = d1.grad(L1, [va, vb])\n",
    "    return dL1_da, dL1_db\n",
    "\n",
    "\n",
    "dva, dvb = run_batched_gradients(d3, d2, d1)\n",
    "print(\"va\", va)\n",
    "print(\"vb\", vb)\n",
    "print(\"dva\", dva)\n",
    "print(\"dvb\", dvb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c66d013e",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "To see that we have done this correctly, we could run the corresponding JAX:\n",
    "\n",
    "```\n",
    "from jax import grad, vmap\n",
    "import jax.numpy as np\n",
    "\n",
    "def simple(a, b):\n",
    "  t = a + b\n",
    "  return t * b\n",
    "\n",
    "def L0(a, b):\n",
    "  return np.sum(simple(a, b))\n",
    "\n",
    "def L1(a, b):\n",
    "  dL0_da, dL0_db = vmap(grad(L0, argnums=(0,1)), in_axes=0)(a, b)\n",
    "  return (dL0_da * dL0_da + dL0_db * dL0_db).sum()\n",
    "\n",
    "va = np.asarray([[0.4556, 0.6323, 0.3489, 0.4017],\n",
    "        [0.0223, 0.1689, 0.2939, 0.5185]])\n",
    "vb = np.asarray([[0.6977, 0.8000, 0.1610, 0.2823],\n",
    "        [0.6816, 0.9152, 0.3971, 0.8742]])\n",
    "dva, dvb = grad(L1, argnums=(0,1))(va, vb)\n",
    "print(\"dva\", dva)\n",
    "print(\"dvb\", dvb)\n",
    "```\n",
    "\n",
    "Looking over the output, the tapes look similar, but we can see that the sizes\n",
    "and the arguments of the operations in question differ (after all, Autograd3 is\n",
    "on the inside of the vmap, while Autograd1 is outside).  But it is still very\n",
    "similar: we could imagine simply varying the dispatcher we use to process backwards\n",
    "depending on when we are executing the tape.  In fact, this is exactly what an\n",
    "initial, non-functorch implementation of PyTorch did to support per-sample\n",
    "gradients.\n",
    "\n",
    "Exercise: modify Autograd.grad to accept a dispatcher, and use that dispatcher\n",
    "instead of self.backward_inner() when running propagator functions.  Then, rewrite\n",
    "the above example so that it only has one level of Autograd:\n",
    "Batched(Autograd(Torch(), create_graph=True)) and show you still get the same\n",
    "result."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2706fca7",
   "metadata": {},
   "source": [
    "OK, so all of this dispatcher business is all nice and explicit, but\n",
    "that's not what JAX/functorch's interface looks like.  How do we\n",
    "bridge the gap?  Our first problem is heving to explicitly thread\n",
    "our Dispatcher object everywhere.  In functorch, we instead implicitly\n",
    "have a \"current mode\" which changes when you enter a grad() or vmap()\n",
    "function.  So let's maintain global current dispatcher, and a way to\n",
    "change what the current dispatcher is.  You can think of this as a\n",
    "singly-linked stack of dispatchers which we push and pop dispatchers\n",
    "onto."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 249,
   "id": "c0d9f0bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "DISPATCHER = Torch()\n",
    "\n",
    "\n",
    "@contextlib.contextmanager\n",
    "def dispatcher(d):\n",
    "    global DISPATCHER\n",
    "    old_d = DISPATCHER\n",
    "    DISPATCHER = d\n",
    "    try:\n",
    "        yield\n",
    "    finally:\n",
    "        DISPATCHER = old_d"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76f9d0e9",
   "metadata": {},
   "source": [
    "A dispatcher mode, however, is not enough.  Remember that in our\n",
    "implementation of Batched, we blindly assumed that all tensors were\n",
    "batched, even if this did not necessarily make sense.  If I have\n",
    "`vmap(lambda bx: bx + y)(x)`, with `x: (B,X)` and `y: (X,)`, the\n",
    "underlying operation should broadcast y to `(B,X)` and then do the\n",
    "addition with x (bx advertises that it has size `(X,)` inside of the\n",
    "vmap'd lambda).  To know this should happen, it is necessary for\n",
    "us to know that y is not a batched tensor, but x is a batched tensor.\n",
    "We'll resolve this with a wrapper class called FuncTensor, which\n",
    "records both the underlying Tensor, as well as the Dispatcher which\n",
    "this tensor is associated with.  In the above example, `bx.dispatcher`\n",
    "might be `Batched(Torch())`, whereas `x.dispatcher` is `Torch()`.\n",
    "\n",
    "So our general strategy is as follows:\n",
    "  1. Every tensor is associated with a dispatcher\n",
    "  2. You can lift tensors to dispatchers which wrap them (which can\n",
    "     trigger some operations, like expand for Batched); this is\n",
    "     implemented by `dispatcher_wraps`\n",
    "  3. To perform an operation between to tensors, lift them so that\n",
    "     they all have the same dispatcher, then do the operation on\n",
    "     that dispatcher."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 250,
   "id": "c80652b4",
   "metadata": {
    "lines_to_end_of_cell_marker": 2,
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "# A dispatcher d1 wraps another dispatcher d2 if d2 is an ancestor of\n",
    "# d1 in the tree structure.  We've defined this relation to be\n",
    "# reflexive, in the same way issubclass(A, A) == True.\n",
    "def dispatcher_wraps(d1, d2):\n",
    "    # Treat this as a reflexive relation\n",
    "    if d1 is d2:\n",
    "        return True\n",
    "    while d1.inner is not None:\n",
    "        d1 = d1.inner\n",
    "        if d1 is d2:\n",
    "            return True\n",
    "    return False\n",
    "\n",
    "\n",
    "# Given a list of arguments, lift them all up to a common dispatcher\n",
    "# level, returning that dispatcher as well as the lifted arguments.\n",
    "# Note that the global current DISPATCHER is also accounted for!\n",
    "# In autodidax, this is `find_top_trace`.\n",
    "def lift_and_unwrap_args(*args):\n",
    "    outermost = DISPATCHER\n",
    "    for a in args:\n",
    "        if dispatcher_wraps(outermost, a.dispatcher):\n",
    "            pass\n",
    "        elif dispatcher_wraps(a.dispatcher, outermost):\n",
    "            # You can make this case an error as well if you don't\n",
    "            # want to support non-lexical functorch tensors\n",
    "            outermost = a.dispatcher\n",
    "        else:\n",
    "            raise TypeError(\"incompatible dispatcher trees\")\n",
    "    return (outermost,) + tuple(a.lift(outermost).tensor for a in args)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9ad2cfe",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "The actual implementation of the wrapper tensor which tracks the\n",
    "Dispatcher for a tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 251,
   "id": "454659ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class FuncTensor:\n",
    "    tensor: Tensor\n",
    "    dispatcher: Dispatcher\n",
    "\n",
    "    # Lift a FuncTensor to an outer dispatcher\n",
    "    def lift(self, d):\n",
    "        # You can only lift to a dispatcher which wraps the dispatcher\n",
    "        # this FuncTensor is associated with (not vice versa, or between\n",
    "        # unrelated FuncTensors).\n",
    "        assert dispatcher_wraps(d, self.dispatcher)\n",
    "        return FuncTensor(d.lift(self.tensor, self.dispatcher), d)\n",
    "\n",
    "    # The general strategy for any operation performed on a tensor, we\n",
    "    # lift all the arguments so that they live on the same dispatcher\n",
    "    # level, and then perform the operation on that dispatcher.  The\n",
    "    # resulting tensor is tagged at whatever dispatcher we had run the\n",
    "    # tensor on.\n",
    "    def __mul__(self, other):\n",
    "        d, self, other = lift_and_unwrap_args(self, other)\n",
    "        return FuncTensor(d.mul(self, other), d)\n",
    "\n",
    "    def __add__(self, other):\n",
    "        d, self, other = lift_and_unwrap_args(self, other)\n",
    "        return FuncTensor(d.add(self, other), d)\n",
    "\n",
    "    def sum(self, dim=None, name=None):\n",
    "        d, self = lift_and_unwrap_args(self)\n",
    "        return FuncTensor(d.sum(self, dim, name=name), d)\n",
    "\n",
    "    def expand(self, sizes):\n",
    "        d, self = lift_and_unwrap_args(self)\n",
    "        return FuncTensor(d.expand(self, sizes), d)\n",
    "\n",
    "    def unsqueeze(self, dim):\n",
    "        d, self = lift_and_unwrap_args(self)\n",
    "        return FuncTensor(d.unsqueeze(self, dim), d)\n",
    "\n",
    "    def squeeze(self, dim):\n",
    "        d, self = lift_and_unwrap_args(self)\n",
    "        return FuncTensor(d.squeeze(self, dim), d)\n",
    "\n",
    "    def size(self):\n",
    "        d, self = lift_and_unwrap_args(self)\n",
    "        return d.size(self)\n",
    "\n",
    "    def dim(self):\n",
    "        d, self = lift_and_unwrap_args(self)\n",
    "        return d.size(self)\n",
    "\n",
    "    # Factory functions like ones do not have any Tensor arguments,\n",
    "    # so they rely solely on the ambient DISPATCHER to determine\n",
    "    # what their semantics should be\n",
    "    @staticmethod\n",
    "    def ones(size):\n",
    "        d = lift_and_unwrap_args()\n",
    "        return d.ones(size)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6063992f",
   "metadata": {},
   "source": [
    "Now we are ready to implement grad.  First, we need some helper\n",
    "functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 252,
   "id": "0926b942",
   "metadata": {},
   "outputs": [],
   "source": [
    "# When we are done doing a vmap/grad, we need to take the results and\n",
    "# lower them back to a lower dispatcher on the stack (this is always\n",
    "# a no-op, in particular, in the vmap case, when we exit vmap the user\n",
    "# gets to see the batched dimension again.)\n",
    "def unlift(t, d):\n",
    "    if isinstance(t, list):\n",
    "        return [unlift(x, d) for x in t]\n",
    "    elif isinstance(t, tuple):\n",
    "        return tuple(unlift(x, d) for x in t)\n",
    "    else:\n",
    "        if t.dispatcher is d:\n",
    "            return t\n",
    "        return unlift(FuncTensor(t.tensor, t.dispatcher.inner), d)\n",
    "\n",
    "\n",
    "# This lets us easily pick out arguments as specified by argnums\n",
    "def filter_argnums(args, argnums):\n",
    "    if isinstance(argnums, int):\n",
    "        return (args[argnums],)\n",
    "    else:\n",
    "        return tuple(args[i] for i in argnums)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cc16dc6",
   "metadata": {},
   "source": [
    "Now grad and vmap!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 253,
   "id": "8778b72f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# For simplicity, these functions only take tuples, not pytrees\n",
    "def grad(f, argnums=0):\n",
    "    @functools.wraps(f)\n",
    "    def wrapped_f(*args):\n",
    "        # We first lift and unwrap all of the arguments which we want\n",
    "        # to pass into the function\n",
    "        old_d, *args = lift_and_unwrap_args(*args)\n",
    "        d = Autograd(old_d)\n",
    "        with dispatcher(d):\n",
    "            # We pass in the functions at the new Autograd level (they\n",
    "            # were lifted to old_d, and lifting to d is a noop)\n",
    "            L = f(*(FuncTensor(a, d) for a in args))\n",
    "            assert L.dispatcher is d\n",
    "            # Run the autograd pass, getting the grads for the inputs\n",
    "            # as specified by argnums\n",
    "            grads = d.grad(L.tensor, filter_argnums(args, argnums))\n",
    "            # Finally, construct the grads at the lower level and return\n",
    "            # them\n",
    "            return [FuncTensor(r, old_d) for r in grads]\n",
    "\n",
    "    return wrapped_f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 254,
   "id": "0134b6c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def vmap(f):\n",
    "    @functools.wraps(f)\n",
    "    def wrapped_f(*args):\n",
    "        # cannot vmap over no arguments as this function uses the\n",
    "        # arguments to determine how large the batch dimension is\n",
    "        # (hypothetically, you could explicitly pass in the batch\n",
    "        # size, and then use this to control factory functions;\n",
    "        # JAX doesn't seem to have a knob to do this)\n",
    "        assert args\n",
    "        old_d, *args = lift_and_unwrap_args(*args)\n",
    "        d = Batched(old_d, length=args[0].size()[0])\n",
    "        for a in args:\n",
    "            assert a.size()[0] == d.length\n",
    "        with dispatcher(d):\n",
    "            # Rewrap all the arguments as batched tensors, then\n",
    "            # unwrap any batched tensors that escape\n",
    "            return unlift(f(*(FuncTensor(a, d) for a in args)), old_d)\n",
    "\n",
    "    return wrapped_f"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "159d3937",
   "metadata": {},
   "source": [
    "Now we can rerun our example using the high level grad/vmap functions!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 255,
   "id": "2a24e765",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Autograd v118: (2, 4) = v79 + v80\n",
      "Autograd v118: (4,) = v79 + v80\n",
      "Autograd v119: (2, 4) = v118 * v80\n",
      "Autograd v119: (4,) = v118 * v80\n",
      "Autograd v120: (2,) = v119.sum(dim=(1,))\n",
      "Autograd v120: () = v119.sum(dim=None)\n",
      "-- Autograd dv120 -------\n",
      "Autograd v122: (2, 1) = v121.unsqueeze(dim=1)\n",
      "Autograd v123: (2, 4) = v122.expand((2, 4))\n",
      "Autograd v124: (2, 4) = v123 * v80\n",
      "Autograd v125: (2, 4) = v123 * v118\n",
      "Autograd v126: (2, 4) = v125 + v124\n",
      "------------------------\n",
      "Autograd v127: (2, 4) = v124 * v124\n",
      "Autograd v128: (2, 4) = v126 * v126\n",
      "Autograd v129: (2, 4) = v127 + v128\n",
      "Autograd v130: () = v129.sum(dim=None)\n",
      "-- Autograd dv130 -------\n",
      "------------------------\n",
      "dva FuncTensor(tensor=tensor([[3.7019, 4.4647, 1.3419, 1.9325],\n",
      "        [2.7711, 3.9985, 2.1762, 4.5337]]), dispatcher=<__main__.Torch object at 0x116031640>)\n",
      "dvb FuncTensor(tensor=tensor([[ 8.7992, 10.5293,  3.0059,  4.4296],\n",
      "        [ 6.9054,  9.8274,  5.1466, 10.8156]]), dispatcher=<__main__.Torch object at 0x116031640>)\n"
     ]
    }
   ],
   "source": [
    "def simple(a, b):\n",
    "    t = a + b\n",
    "    return t * b\n",
    "\n",
    "\n",
    "def L0(a, b):\n",
    "    return simple(a, b).sum()\n",
    "\n",
    "\n",
    "def L1(a, b):\n",
    "    dL0_da, dL0_db = vmap(grad(L0, argnums=(0, 1)))(a, b)\n",
    "    return (dL0_da * dL0_da + dL0_db * dL0_db).sum()\n",
    "\n",
    "\n",
    "fva = FuncTensor(va, DISPATCHER)\n",
    "fvb = FuncTensor(vb, DISPATCHER)\n",
    "dva, dvb = grad(L1, argnums=(0, 1))(fva, fvb)\n",
    "print(\"dva\", dva)\n",
    "print(\"dvb\", dvb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "268470a8",
   "metadata": {},
   "source": [
    "Because FuncTensors are associated with the ambient dispatcher they\n",
    "were created from, they are also allowed to escape from the context in\n",
    "which they were defined, allowing for non-lexical, imperative\n",
    "transform API.  For example, batching over module parameters is\n",
    "problematic today, but all we need to do is tweak the FuncTensor's\n",
    "dispatchers appropriately and everything works out."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 256,
   "id": "29d937c8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "expect tensor([[ 1.5606, -1.8643,  0.4406],\n",
      "        [ 1.0271,  0.9861,  0.5888]])\n",
      "output tensor([[ 1.5606, -1.8643,  0.4406],\n",
      "        [ 1.0271,  0.9861,  0.5888]])\n"
     ]
    }
   ],
   "source": [
    "\n",
    "PlainTensor = lambda t: FuncTensor(torch.randn(N), DISPATCHER)\n",
    "BatchedTensor = lambda t: FuncTensor(t, Batched(DISPATCHER, length=B))\n",
    "\n",
    "class ScaleBiasModule:\n",
    "    weight: FuncTensor\n",
    "    bias: FuncTensor\n",
    "\n",
    "    def __init__(self, N):\n",
    "        self.weight = PlainTensor(torch.randn(N))\n",
    "        self.bias = PlainTensor(torch.randn(N))\n",
    "\n",
    "    def forward(self, input):\n",
    "        return self.weight * input + self.bias\n",
    "\n",
    "\n",
    "B = 2\n",
    "N = 3\n",
    "m = ScaleBiasModule(N)\n",
    "# Ensemble weights only; input is not batched\n",
    "m.weight = BatchedTensor(torch.randn(B, N))\n",
    "input = PlainTensor(torch.randn(N))\n",
    "output = m.forward(input)\n",
    "print(\n",
    "    \"expect\", input.tensor.unsqueeze(0) * m.weight.tensor + m.bias.tensor.unsqueeze(0)\n",
    ")\n",
    "print(\"output\", output.tensor)"
   ]
  }
 ],
 "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
}