rocm_jax/docs/autodidax.ipynb

{
 "cells": [
  {
   "cell_type": "raw",
   "id": "surrounded-agent",
   "metadata": {},
   "source": [
    "---\n",
    "Copyright 2021 Google LLC\n",
    "\n",
    "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "you may not use this file except in compliance with the License.\n",
    "You may obtain a copy of the License at\n",
    "\n",
    "    https://www.apache.org/licenses/LICENSE-2.0\n",
    "\n",
    "Unless required by applicable law or agreed to in writing, software\n",
    "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "See the License for the specific language governing permissions and\n",
    "limitations under the License.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Autodidax: JAX core from scratch\n",
    "\n",
    "Ever want to learn how JAX works, but the implementation seemed impenetrable?\n",
    "Well, you're in luck! By reading this tutorial, you'll learn every big idea in\n",
    "JAX's core system. You'll even get clued into our weird jargon!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 1: Transformations as interpreters: standard evaluation, `jvp`, and `vmap`\n",
    "\n",
    "We want to transform functions that look like this:\n",
    "\n",
    "```python\n",
    "def f(x):\n",
    "  y = sin(x) * 2.\n",
    "  z = - y + x\n",
    "  return z\n",
    "```\n",
    "\n",
    "Think of functions like `sin` and the arithmetic operations underlying the\n",
    "infix operators (`mul`, `add`, and `neg`) as primitive operations, meaning\n",
    "atomic units of processing rather than compositions.\n",
    "\n",
    "\"Transform\" means \"interpret differently.\" Instead of standard interpretation\n",
    "where we apply primitive operations to numerical inputs to produce numerical\n",
    "outputs, we want to override primitive application and let different values\n",
    "flow through our program. For example, we might want to replace the\n",
    "application of every primitive with an application of [its JVP\n",
    "rule](https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html),\n",
    "and let primal-tangent pairs flow through our program. Moreover, we want to be\n",
    "able to comopse multiple transformations, leading to stacks of interpreters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "### JAX core machinery\n",
    "\n",
    "We can implement stacks of interpreters and even have them all discharge on\n",
    "the fly as we execute the Python function to be transformed. To start, let's\n",
    "define these primitives so that we can intercept their application:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_end_of_cell_marker": 0,
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "from typing import NamedTuple\n",
    "\n",
    "class Primitive(NamedTuple):\n",
    "  name: str\n",
    "\n",
    "add_p = Primitive('add')\n",
    "mul_p = Primitive('mul')\n",
    "neg_p = Primitive(\"neg\")\n",
    "sin_p = Primitive(\"sin\")\n",
    "cos_p = Primitive(\"cos\")\n",
    "reduce_sum_p = Primitive(\"reduce_sum\")\n",
    "greater_p = Primitive(\"greater\")\n",
    "transpose_p = Primitive(\"transpose\")\n",
    "broadcast_p = Primitive(\"broadcast\")\n",
    "\n",
    "def add(x, y): return bind1(add_p, x, y)\n",
    "def mul(x, y): return bind1(mul_p, x, y)\n",
    "def neg(x): return bind1(neg_p, x)\n",
    "def sin(x): return bind1(sin_p, x)\n",
    "def cos(x): return bind1(cos_p, x)\n",
    "def reduce_sum(x, axis=None): return bind1(reduce_sum_p, x, axis=axis)\n",
    "def greater(x, y): return bind1(greater_p, x, y)\n",
    "def transpose(x, perm): return bind1(transpose_p, perm=perm)\n",
    "def broadcast(x, shape, axes): return bind1(broadcast_p, x, shape=shape, axes=axes)\n",
    "\n",
    "def bind1(prim, *args, **params):\n",
    "  out, = bind(prim, *args, **params)\n",
    "  return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll set up array data types and infix operator methods in a moment.\n",
    "\n",
    "A `Primitive` is just an object with a name, to which we attach our\n",
    "interpretation rules (one for each transformation). The `bind` function is our\n",
    "interception point: it'll figure out which transformation rule to apply, based\n",
    "on how the arguments are boxed in tracers and what interpreters are active.\n",
    "\n",
    "The functions that user code calls, like `add` and `sin`, are just wrappers\n",
    "around calls to `bind`. These wrappers let us control how arguments are passed\n",
    "to `bind`, and in particular we follow a handy internal convention: when we\n",
    "call `bind`, we pass values representing array data as positional arguments,\n",
    "and we pass metadata like the `axis` argument to `sum_p` via keyword. This\n",
    "calling convention simplifies some core logic (since e.g. instances of the\n",
    "`Tracer` class to be defined below can only occurr in positional arguments to\n",
    "`bind`). The wrappers can also provide docstrings!\n",
    "\n",
    "We represent active interpreters as a stack. The stack is just a simple\n",
    "`list`, and each element is a container with an integer level (corresponding\n",
    "to the element's height in the stack), an interpreter type (which we'll call a\n",
    "`trace_type`), and an optional field for any global data the interpreter\n",
    "needs. We call each element a `MainTrace`, though maybe \"Interpreter\" would be\n",
    "more descriptive."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "from contextlib import contextmanager\n",
    "from typing import Type, List, Optional, Any"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "class MainTrace(NamedTuple):\n",
    "  level: int\n",
    "  trace_type: Type['Trace']\n",
    "  global_data: Optional[Any]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "trace_stack: List[MainTrace] = []\n",
    "dynamic_trace: Optional[MainTrace] = None  # to be employed in Part 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "@contextmanager\n",
    "def new_main(trace_type: Type['Trace'], global_data=None):\n",
    "  level = len(trace_stack)\n",
    "  main = MainTrace(level, trace_type, global_data)\n",
    "  trace_stack.append(main)\n",
    "\n",
    "  try:\n",
    "    yield main\n",
    "  finally:\n",
    "    trace_stack.pop()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When we're about to apply a transformation, we'll push another interpreter\n",
    "onto the stack using `new_main`. Then, as we apply primitives in the function,\n",
    "we can think of the `bind` first being interpreted by the trace at the top of\n",
    "the stack (i.e. with the highest level). If that first interpreter itself\n",
    "binds other primitives in its interpretation rule for the primitive, like how\n",
    "the JVP rule of `sin_p` might bind `cos_p` and `mul_p`, then those `bind`\n",
    "calls will be handled by the interpreter at the next level down.\n",
    "\n",
    "What goes at the bottom of the interpreter stack? At the bottom, we know all\n",
    "the transformation interpreters are finished, and we just want to do standard\n",
    "evaluation. So at the bottom we'll put an evaluation interpreter.\n",
    "\n",
    "Let's sketch out the interface for interpreters, which is based on the `Trace`\n",
    "and `Tracer` base classes. A `Tracer` represents a boxed-up value, perhaps\n",
    "carrying some extra context data used by the interpreter. A `Trace` handles\n",
    "boxing up vales into `Tracers` and also handles primitive application."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "class Trace:\n",
    "  main: MainTrace\n",
    "\n",
    "  def __init__(self, main: MainTrace) -> None:\n",
    "    self.main = main\n",
    "\n",
    "  def pure(self, val): assert False  # must override\n",
    "  def lift(self, val): assert False  # must override\n",
    "\n",
    "  def process_primitive(self, primitive, tracers, params):\n",
    "    assert False  # must override"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first two methods are about boxing up values in `Tracer`s, which are the\n",
    "objects that flow through the Python programs we transform. The last method is\n",
    "the callback we'll use to interpret primitive application.\n",
    "\n",
    "The `Trace` itself doesn't contain any data, other than a reference to its\n",
    "corresponding `MainTrace` instance. In fact, multiple instances of a `Trace`\n",
    "might be created and discarded during an application of a transformation,\n",
    "whereas only a single `MainTrace` instance is created per application of a\n",
    "transformation.\n",
    "\n",
    "As for `Tracer`s themselves, each one carries an abstract value (and forwards\n",
    "infix operators to it), and the rest is up to the transformation. (The\n",
    "relationship between `Tracer`s and `AbstractValue`s is that there's one\n",
    "`Tracer` per transformation, and at least one `AbstractValue` per base type,\n",
    "like arrays.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from typing import Tuple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "class Tracer:\n",
    "  _trace: Trace\n",
    "\n",
    "  __array_priority__ = 1000\n",
    "\n",
    "  @property\n",
    "  def aval(self):\n",
    "    assert False  # must override\n",
    "\n",
    "  def full_lower(self):\n",
    "    return self  # default implementation\n",
    "\n",
    "  def __neg__(self): return self.aval._neg(self)\n",
    "  def __add__(self, other): return self.aval._add(self, other)\n",
    "  def __radd__(self, other): return self.aval._radd(self, other)\n",
    "  def __mul__(self, other): return self.aval._mul(self, other)\n",
    "  def __rmul__(self, other): return self.aval._rmul(self, other)\n",
    "  def __gt__(self, other): return self.aval._gt(self, other)\n",
    "  def __bool__(self): return self.aval._bool(self)\n",
    "  def __nonzero__(self): return self.aval._nonzero(self)\n",
    "\n",
    "  def __getattr__(self, name):\n",
    "    try:\n",
    "      return getattr(self.aval, name)\n",
    "    except AttributeError:\n",
    "      raise AttributeError(f\"{self.__class__.__name__} has no attribute {name}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "class ShapedArray:\n",
    "  array_abstraction_level = 1\n",
    "  shape: Tuple[int]\n",
    "  dtype: np.dtype\n",
    "\n",
    "  def __init__(self, shape, dtype):\n",
    "    self.shape = shape\n",
    "    self.dtype = dtype\n",
    "\n",
    "  @property\n",
    "  def ndim(self):\n",
    "    return len(self.shape)\n",
    "\n",
    "  _neg = staticmethod(neg)\n",
    "  _add = staticmethod(add)\n",
    "  _radd = staticmethod(add)\n",
    "  _mul = staticmethod(mul)\n",
    "  _rmul = staticmethod(mul)\n",
    "  _gt = staticmethod(greater)\n",
    "\n",
    "  @staticmethod\n",
    "  def _bool(tracer):\n",
    "    raise Exception(\"ShapedArray can't be unambiguously converted to bool\")\n",
    "\n",
    "  @staticmethod\n",
    "  def _nonzero(tracer):\n",
    "    raise Exception(\"ShapedArray can't be unambiguously converted to bool\")\n",
    "\n",
    "  def str_short(self):\n",
    "    return f'{self.dtype.name}[{\",\".join(str(d) for d in self.shape)}]'\n",
    "\n",
    "  def __hash__(self):\n",
    "    return hash((self.shape, self.dtype))\n",
    "\n",
    "  def __eq__(self, other):\n",
    "    return (type(self) is type(other) and\n",
    "            self.shape == other.shape and self.dtype == other.dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "class ConcreteArray(ShapedArray):\n",
    "  array_abstraction_level = 2\n",
    "  val: np.ndarray\n",
    "\n",
    "  def __init__(self, val):\n",
    "    self.val = val\n",
    "    self.shape = val.shape\n",
    "    self.dtype = val.dtype\n",
    "\n",
    "  @staticmethod\n",
    "  def _bool(tracer):\n",
    "    return bool(tracer.aval.val)\n",
    "\n",
    "  @staticmethod\n",
    "  def _nonzero(tracer):\n",
    "    return bool(tracer.aval.val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_end_of_cell_marker": 0,
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def get_aval(x):\n",
    "  if isinstance(x, Tracer):\n",
    "    return x.aval\n",
    "  else:\n",
    "    return ConcreteArray(np.asarray(x))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that we actually have two `AbstractValue`s for arrays, representing\n",
    "different levels of abstraction. A `ShapedArray` represents the set of all\n",
    "possible arrays with a given shape and dtype. A `ConcreteArray` represents a\n",
    "singleton set consisting of a single array value.\n",
    "\n",
    "Now that we've set up the interpreter stack, the Trace/Tracer API for\n",
    "interpreters, and abstract values, we can come back to implement `bind`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def bind(prim, *args, **params):\n",
    "  top_trace = find_top_trace(args)\n",
    "  tracers = [full_raise(top_trace, arg) for arg in args]\n",
    "  outs = top_trace.process_primitive(prim, tracers, params)\n",
    "  return [full_lower(out) for out in outs]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The main action is that we call `find_top_trace` to figure out which\n",
    "interpreter should handle this primitive application. We then call that top\n",
    "trace's `process_primitive` so that the trace can apply its interpretation\n",
    "rule. The calls to `full_raise` just ensure that the inputs are boxed in the\n",
    "top trace's `Tracer` instances, and the call to `full_lower` is an optional\n",
    "optimization so that we unbox values out of `Tracer`s as much as possible."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "from operator import attrgetter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def find_top_trace(xs) -> Trace:\n",
    "  top_main = max((x._trace.main for x in xs if isinstance(x, Tracer)),\n",
    "                 default=trace_stack[0], key=attrgetter('level'))\n",
    "  if dynamic_trace and dynamic_trace.level > top_main.level:\n",
    "    top_main = dynamic_trace\n",
    "  return top_main.trace_type(top_main)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In words, ignoring the `dynamic_trace` step until Part 3, `find_top_trace`\n",
    "returns the highest-level interpreter associated with the `Tracer`s on its\n",
    "inputs, and otherwise returns the interpreter at the bottom of the stack\n",
    "(which is always an evaluation trace, at least for now). This is a deviation\n",
    "from the description above, where we always start by running the interpreter\n",
    "at the top of the stack and then work our way down, applying every interpreter\n",
    "in the stack. Instead, we're only applying an interpreter when the input\n",
    "arguments to a primitive bind are boxed in a `Tracer` corresponding to that\n",
    "interpreter. This optimization lets us skip irrelevant transformations, but\n",
    "bakes in an assumption that transformations mostly follow data dependence\n",
    "(except for the special bottom-of-the-stack interpreter, which interprets\n",
    "everything).\n",
    "\n",
    "An alternative would be to have every interpreter in the stack interpret every\n",
    "operation. That's worth exploring! JAX is designed around data dependence in\n",
    "large part because that's so natural for automatic differentiation, and JAX's\n",
    "roots are in autodiff. But it may be over-fit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def full_lower(val: Any):\n",
    "  if isinstance(val, Tracer):\n",
    "    return val.full_lower()\n",
    "  else:\n",
    "    return val"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def full_raise(trace: Trace, val: Any) -> Tracer:\n",
    "  if not isinstance(val, Tracer):\n",
    "    return trace.pure(val)\n",
    "  level = trace.main.level\n",
    "  if val._trace.main is trace.main:\n",
    "    return val\n",
    "  elif val._trace.main.level < level:\n",
    "    return trace.lift(val)\n",
    "  elif val._trace.main.level > level:\n",
    "    raise Exception(f\"Can't lift level {val._trace.main.level} to {level}.\")\n",
    "  else:  # val._trace.level == level\n",
    "    raise Exception(f\"Different traces at same level: {val._trace}, {trace}.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The logic in `full_raise` serves to box values into `Tracer`s for a particular\n",
    "`Trace`, calling different methods on the `Trace` based on context:\n",
    "`Trace.pure` is called on non-`Tracer` constants, and `Trace.lift` is called\n",
    "for values that are already `Tracer`s from a lower-level interpreter. These\n",
    "two methods could share the same implementation, but by distinguishing them in\n",
    "the core logic we can provide more information to the `Trace` subclass.\n",
    "\n",
    "That's it for the JAX core! Now we can start adding interpreters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluation interpreter\n",
    "\n",
    "We'll start with the simplest interpreter: the evaluation interpreter that\n",
    "will sit at the bottom of the interpreter stack."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "class EvalTrace(Trace):\n",
    "  pure = lift = lambda self, x: x  # no boxing in Tracers needed\n",
    "\n",
    "  def process_primitive(self, primitive, tracers, params):\n",
    "    return impl_rules[primitive](*tracers, **params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trace_stack.append(MainTrace(0, EvalTrace, None))  # special bottom of the stack"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "impl_rules = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "impl_rules[add_p] = lambda x, y: [np.add(x, y)]\n",
    "impl_rules[mul_p] = lambda x, y: [np.multiply(x, y)]\n",
    "impl_rules[neg_p] = lambda x: [np.negative(x)]\n",
    "impl_rules[sin_p] = lambda x: [np.sin(x)]\n",
    "impl_rules[cos_p] = lambda x: [np.cos(x)]\n",
    "impl_rules[reduce_sum_p] = lambda x, *, axis: [np.sum(x, axis)]\n",
    "impl_rules[greater_p] = lambda x, y: [np.greater(x, y)]\n",
    "impl_rules[transpose_p] = lambda x, *, perm: [np.transpose(x, perm)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def broadcast_impl(x, *, shape, axes):\n",
    "  return [np.broadcast_to(np.expand_dims(x, axes), shape)]\n",
    "impl_rules[broadcast_p] = broadcast_impl"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With this interpreter, we can evaluate user functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def f(x):\n",
    "  y = sin(x) * 2.\n",
    "  z = - y + x\n",
    "  return z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.7177599838802657\n"
     ]
    }
   ],
   "source": [
    "print(f(3.0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Woo! Like going around in a big circle. But the point of this indirection is\n",
    "that now we can add some real transformations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Forward-mode autodiff with `jvp`\n",
    "\n",
    "First, a few helper functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def zeros_like(val):\n",
    "  return np.zeros_like(val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def unzip2(pairs):\n",
    "  lst1, lst2 = [], []\n",
    "  for x1, x2 in pairs:\n",
    "    lst1.append(x1)\n",
    "    lst2.append(x2)\n",
    "  return lst1, lst2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "map_ = map\n",
    "def map(f, *xs):\n",
    "  return list(map_(f, *xs))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `Tracer` for forward-mode autodiff carries a primal-tangent pair. The\n",
    "`Trace` applies JVP rules."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "class JVPTracer(Tracer):\n",
    "  def __init__(self, trace, primal, tangent):\n",
    "    self._trace = trace\n",
    "    self.primal = primal\n",
    "    self.tangent = tangent\n",
    "\n",
    "  @property\n",
    "  def aval(self):\n",
    "    return get_aval(self.primal)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "class JVPTrace(Trace):\n",
    "  pure = lift = lambda self, val: JVPTracer(self, val, zeros_like(val))\n",
    "\n",
    "  def process_primitive(self, primitive, tracers, params):\n",
    "    primals_in, tangents_in = unzip2((t.primal, t.tangent) for t in tracers)\n",
    "    jvp_rule = jvp_rules[primitive]\n",
    "    primal_outs, tangent_outs = jvp_rule(primals_in, tangents_in, **params)\n",
    "    return [JVPTracer(self, x, t) for x, t in zip(primal_outs, tangent_outs)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "jvp_rules = {}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice both `lift` and `sublift` package a value into a `JVPTracer` with the\n",
    "minimal amount of context, which is a zero tangent value.\n",
    "\n",
    "Let's add some JVP rules for primitives:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def add_jvp(primals, tangents):\n",
    "  (x, y), (x_dot, y_dot) = primals, tangents\n",
    "  return [x + y], [x_dot + y_dot]\n",
    "jvp_rules[add_p] = add_jvp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def mul_jvp(primals, tangents):\n",
    "  (x, y), (x_dot, y_dot) = primals, tangents\n",
    "  return [x * y], [x_dot * y + x * y_dot]\n",
    "jvp_rules[mul_p] = mul_jvp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def sin_jvp(primals, tangents):\n",
    "  (x,), (x_dot,) = primals, tangents\n",
    "  return [sin(x)], [cos(x) * x_dot]\n",
    "jvp_rules[sin_p] = sin_jvp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def cos_jvp(primals, tangents):\n",
    "  (x,), (x_dot,) = primals, tangents\n",
    "  return [cos(x)], [-sin(x) * x_dot]\n",
    "jvp_rules[cos_p] = cos_jvp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def neg_jvp(primals, tangents):\n",
    "  (x,), (x_dot,) = primals, tangents\n",
    "  return [neg(x)], [neg(x_dot)]\n",
    "jvp_rules[neg_p] = neg_jvp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def reduce_sum_jvp(primals, tangents, *, axis):\n",
    "  (x,), (x_dot,) = primals, tangents\n",
    "  return [reduce_sum(x, axis)], [reduce_sum(x_dot, axis)]\n",
    "jvp_rules[reduce_sum_p] = reduce_sum_jvp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def greater_jvp(primals, tangents):\n",
    "  (x, y), _ = primals, tangents\n",
    "  out_primal = greater(x, y)\n",
    "  return [out_primal], [zeros_like(out_primal)]\n",
    "jvp_rules[greater_p] = greater_jvp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we add a transformation API to kick off the trace:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def jvp_v1(f, primals, tangents):\n",
    "  with new_main(JVPTrace) as main:\n",
    "    trace = JVPTrace(main)\n",
    "    tracers_in = [JVPTracer(trace, x, t) for x, t in zip(primals, tangents)]\n",
    "    out = f(*tracers_in)\n",
    "    tracer_out = full_raise(trace, out)\n",
    "    primal_out, tangent_out = tracer_out.primal, tracer_out.tangent\n",
    "  return primal_out, tangent_out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And with that, we can differentiate!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-0.9899924966004454\n",
      "-0.9899924966004454\n"
     ]
    }
   ],
   "source": [
    "x = 3.0\n",
    "y, sin_deriv_at_3 = jvp_v1(sin, (x,), (1.0,))\n",
    "print(sin_deriv_at_3)\n",
    "print(cos(3.0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.7177599838802657\n",
      "2.979984993200891\n"
     ]
    }
   ],
   "source": [
    "def f(x):\n",
    "  y = sin(x) * 2.\n",
    "  z = - y + x\n",
    "  return z\n",
    "\n",
    "x, xdot = 3., 1.\n",
    "y, ydot = jvp_v1(f, (x,), (xdot,))\n",
    "print(y)\n",
    "print(ydot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-0.9899924966004454\n",
      "-0.1411200080598672\n",
      "0.9899924966004454\n",
      "0.1411200080598672\n"
     ]
    }
   ],
   "source": [
    "def deriv(f):\n",
    "  return lambda x: jvp_v1(f, (x,), (1.,))[1]\n",
    "\n",
    "print(deriv(sin)(3.))\n",
    "print(deriv(deriv(sin))(3.))\n",
    "print(deriv(deriv(deriv(sin)))(3.))\n",
    "print(deriv(deriv(deriv(deriv(sin))))(3.))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_end_of_cell_marker": 0,
    "lines_to_next_cell": 1
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.0\n",
      "1.0\n"
     ]
    }
   ],
   "source": [
    "def f(x):\n",
    "  if x > 0.:  # Python control flow\n",
    "    return 2. * x\n",
    "  else:\n",
    "    return x\n",
    "\n",
    "print(deriv(f)(3.))\n",
    "print(deriv(f)(-3.))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pytrees and flattening user functions' inputs and outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A limitation with `jvp_v1` is that it assumes the user function accepts arrays\n",
    "as positional arguments and produces a single array as output. What if it\n",
    "produced a list as output? Or accepted nested containers as inputs? It would\n",
    "be a pain to deal with all the possible containers in inputs and outputs at\n",
    "every layer of the stack. Instead, we can wrap the user function so that the\n",
    "wrapped version accepts arrays as inputs and returns a flat list of arrays as\n",
    "output. The wrapper just needs to unflatten its input, call the user function,\n",
    "and flatten the output.\n",
    "\n",
    "Here's how we'd like to write `jvp`, assuming the user always gives us\n",
    "functions that take arrays as inputs and produces a flat list of arrays as\n",
    "outputs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def jvp_flat(f, primals, tangents):\n",
    "  with new_main(JVPTrace) as main:\n",
    "    trace = JVPTrace(main)\n",
    "    tracers_in = [JVPTracer(trace, x, t) for x, t in zip(primals, tangents)]\n",
    "    outs = f(*tracers_in)\n",
    "    tracers_out = [full_raise(trace, out) for out in outs]\n",
    "    primals_out, tangents_out = unzip2((t.primal, t.tangent) for t in tracers_out)\n",
    "  return primals_out, tangents_out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To support user functions that have arbitrary containers in the inputs and\n",
    "outputs, here's how we'd write the user-facing `jvp` wrapper:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def jvp(f, primals, tangents):\n",
    "  primals_flat, in_tree = tree_flatten(primals)\n",
    "  tangents_flat, in_tree2 = tree_flatten(tangents)\n",
    "  if in_tree != in_tree2: raise TypeError\n",
    "  f, out_tree = flatten_fun(f, in_tree)\n",
    "  primals_out_flat, tangents_out_flat = jvp_flat(f, primals_flat, tangents_flat)\n",
    "  primals_out = tree_unflatten(out_tree(), primals_out_flat)\n",
    "  tangents_out = tree_unflatten(out_tree(), tangents_out_flat)\n",
    "  return primals_out, tangents_out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that we had to plumb the tree structure of the user function output\n",
    "back to the caller of `flatten_fun`. That information isn't available until we\n",
    "actually run the user function, so `flatten_fun` just returns a reference to a\n",
    "mutable cell, represented as a thunk. These side-effects are safe because we\n",
    "always run the user function exactly once. (This safe regime is the reason for\n",
    "the \"linear\" name in `linear_util.py`, in the sense of [linear\n",
    "types](https://en.wikipedia.org/wiki/Substructural_type_system).)\n",
    "\n",
    "All that remains is to write `tree_flatten`, `tree_unflatten`, and\n",
    "`flatten_fun`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def flatten_fun(f, in_tree):\n",
    "  store = Store()\n",
    "\n",
    "  def flat_fun(*args_flat):\n",
    "    pytree_args = tree_unflatten(in_tree, args_flat)\n",
    "    out = f(*pytree_args)\n",
    "    out_flat, out_tree = tree_flatten(out)\n",
    "    store.set_value(out_tree)\n",
    "    return out_flat\n",
    "\n",
    "  return flat_fun, store"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "class Empty: pass\n",
    "empty = Empty()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "class Store:\n",
    "  val = empty\n",
    "\n",
    "  def set_value(self, val):\n",
    "    assert self.val is empty\n",
    "    self.val = val\n",
    "\n",
    "  def __call__(self):\n",
    "    return self.val"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_end_of_cell_marker": 0,
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "import itertools as it\n",
    "from typing import Callable, Type, Hashable, Dict, Iterable, Iterator\n",
    "\n",
    "class NodeType(NamedTuple):\n",
    "  to_iterable: Callable\n",
    "  from_iterable: Callable\n",
    "\n",
    "node_types: Dict[Type, NodeType] = {\n",
    "    tuple: NodeType(lambda t: (None, t), lambda _, xs: tuple(xs)),\n",
    "    list: NodeType( lambda l: (None, l), lambda _, xs:  list(xs)),\n",
    "    dict: NodeType(lambda d: map(tuple, unzip2(sorted(d.items()))),\n",
    "                   lambda keys, vals: dict(zip(keys, vals))),\n",
    "}\n",
    "\n",
    "class PyTreeDef(NamedTuple):\n",
    "  node_type: NodeType\n",
    "  node_metadata: Hashable\n",
    "  child_treedefs: Tuple['PyTreeDef']\n",
    "\n",
    "class Leaf: pass\n",
    "leaf = Leaf()\n",
    "\n",
    "def tree_flatten(x: Any) -> Tuple[List[Any], PyTreeDef]:\n",
    "  children_iter, treedef = _tree_flatten(x)\n",
    "  return list(children_iter), treedef\n",
    "\n",
    "def _tree_flatten(x: Any) -> Tuple[Iterable, PyTreeDef]:\n",
    "  node_type = node_types.get(type(x))\n",
    "  if node_type:\n",
    "    node_metadata, children = node_type.to_iterable(x)\n",
    "    children_flat, child_trees = unzip2(map(_tree_flatten, children))\n",
    "    flattened = it.chain.from_iterable(children_flat)\n",
    "    return flattened, PyTreeDef(node_type, node_metadata, tuple(child_trees))\n",
    "  else:\n",
    "    return [x], leaf\n",
    "\n",
    "def tree_unflatten(treedef: PyTreeDef, xs: List[Any]) -> Any:\n",
    "  return _tree_unflatten(treedef, iter(xs))\n",
    "\n",
    "def _tree_unflatten(treedef: PyTreeDef, xs: Iterator) -> Any:\n",
    "  if treedef is leaf:\n",
    "    return next(xs)\n",
    "  else:\n",
    "    children = (_tree_unflatten(t, xs) for t in treedef.child_treedefs)\n",
    "    return treedef.node_type.from_iterable(treedef.node_metadata, children)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With this pytree-handling `jvp` impelmentation, we can now handle arbitrary\n",
    "input and output containers. That'll come in handy with future transformations\n",
    "too!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "def f(x):\n",
    "  y = sin(x) * 2.\n",
    "  z = - y + x\n",
    "  return {'hi': z, 'there': [x, y]}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "x, xdot = 3., 1.\n",
    "y, ydot = jvp(f, (x,), (xdot,))\n",
    "print(y)\n",
    "print(ydot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Vectorized batching with `vmap`\n",
    "\n",
    "First, a couple helper functions, one for producing mapped abstract values\n",
    "from unmapped ones (by removing an axis), and one for moving batch dimensions\n",
    "around:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def mapped_aval(batch_dim, aval):\n",
    "  shape = list(aval.shape)\n",
    "  del shape[batch_dim]\n",
    "  return ShapedArray(tuple(shape), aval.dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def move_batch_axis(axis_size, src, dst, x):\n",
    "  if src is not_mapped:\n",
    "    target_shape = list(np.shape(x))\n",
    "    target_shape.insert(dst, axis_size)\n",
    "    return broadcast(x, target_shape, [dst])\n",
    "  elif src == dst:\n",
    "    return x\n",
    "  else:\n",
    "    return moveaxis(x, src, dst)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def moveaxis(x, src: int, dst: int):\n",
    "  perm = [i for i in range(np.ndim(x)) if i != src]\n",
    "  perm.insert(dst, src)\n",
    "  return transpose(x, perm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `Tracer` for vectorized batching carries a batched value and an optional\n",
    "integer indicating which axis (if any) is the batch axis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "from typing import Union"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "class NotMapped: pass\n",
    "not_mapped = NotMapped()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "BatchAxis = Union[NotMapped, int]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "class BatchTracer(Tracer):\n",
    "  def __init__(self, trace, val, batch_dim: BatchAxis):\n",
    "    self._trace = trace\n",
    "    self.val = val\n",
    "    self.batch_dim = batch_dim\n",
    "\n",
    "  @property\n",
    "  def aval(self):\n",
    "    if self.batch_dim is not_mapped:\n",
    "      return get_aval(self.val)\n",
    "    else:\n",
    "      return mapped_aval(self.batch_dim, get_aval(self.val))\n",
    "\n",
    "  def full_lower(self):\n",
    "    if self.batch_dim is not_mapped:\n",
    "      return full_lower(self.val)\n",
    "    else:\n",
    "      return self"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "class BatchTrace(Trace):\n",
    "  pure = lift = lambda self, val: BatchTracer(self, val, not_mapped)\n",
    "\n",
    "  def process_primitive(self, primitive, tracers, params):\n",
    "    vals_in, bdims_in = unzip2((t.val, t.batch_dim) for t in tracers)\n",
    "    vmap_rule = vmap_rules[primitive]\n",
    "    val_outs, bdim_outs = vmap_rule(self.axis_size, vals_in, bdims_in, **params)\n",
    "    return [BatchTracer(self, x, bd) for x, bd in zip(val_outs, bdim_outs)]\n",
    "\n",
    "  @property\n",
    "  def axis_size(self):\n",
    "    return self.main.global_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vmap_rules = {}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we've implemented the optional `Tracer.full_lower` method, which lets us\n",
    "peel off a batching tracer if it's not needed because it doesn't represent a\n",
    "batched value.\n",
    "\n",
    "For `BatchTrace`, analogous to `JVPTrace`, the methods `pure` and `lift` just\n",
    "box a value in a `BatchTracer` with the minimal amount of context, which in\n",
    "this case is a `batch_dim` taking the sentinel value `not_mapped`. Notice we\n",
    "use the `MainTrace`'s interpreter-global data field to store the batch axis\n",
    "size.\n",
    "\n",
    "Next we can define batching interpreter rules for each primitive:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "from functools import partial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def broadcasting_binop_batching_rule(op, axis_size, vals_in, dims_in):\n",
    "  (x, y), (x_bdim, y_bdim) = vals_in, dims_in\n",
    "  if x_bdim != y_bdim:\n",
    "    y = move_batch_axis(axis_size, y_bdim, x_bdim, y)\n",
    "  return [op(x, y)], [x_bdim]\n",
    "vmap_rules[add_p] = partial(broadcasting_binop_batching_rule, add)\n",
    "vmap_rules[mul_p] = partial(broadcasting_binop_batching_rule, mul)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def vectorized_unop_batching_rule(op, axis_size, vals_in, dims_in):\n",
    "  (x,), (x_bdim,) = vals_in, dims_in\n",
    "  return [op(x)], [x_bdim]\n",
    "vmap_rules[sin_p] = partial(vectorized_unop_batching_rule, sin)\n",
    "vmap_rules[cos_p] = partial(vectorized_unop_batching_rule, cos)\n",
    "vmap_rules[neg_p] = partial(vectorized_unop_batching_rule, neg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def reduce_sum_batching_rule(axis_size, vals_in, dims_in, *, axis):\n",
    "  (x,), (x_bdim,) = vals_in, dims_in\n",
    "  new_axis = axis + (x_bdim <= axis)\n",
    "  out_bdim = x_bdim - (new_axis < x_bdim)\n",
    "  return [reduce_sum(x, new_axis)], [out_bdim]\n",
    "vmap_rules[reduce_sum_p] = reduce_sum_batching_rule"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "-"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we add a transformation API to kick off the trace:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def vmap_flat(f, in_axes, *args):\n",
    "  axis_size, = {x.shape[ax] for x, ax in zip(args, in_axes)\n",
    "                if ax is not not_mapped}\n",
    "  with new_main(BatchTrace, axis_size) as main:\n",
    "    trace = BatchTrace(main)\n",
    "    tracers_in = [BatchTracer(trace, x, ax) if ax is not None else x\n",
    "                  for x, ax in zip(args, in_axes)]\n",
    "    outs = f(*tracers_in)\n",
    "    tracers_out = [full_raise(trace, out) for out in outs]\n",
    "    vals_out, bdims_out = unzip2((t.val, t.batch_dim) for t in tracers_out)\n",
    "  outs_transposed = [move_batch_axis(axis_size, bdim, 0, val_out)\n",
    "                     for val_out, bdim in zip(vals_out, bdims_out)]\n",
    "  return outs_transposed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def vmap(f, in_axes):\n",
    "  def batched_f(*args):\n",
    "    args_flat, in_tree = tree_flatten(args)\n",
    "    in_axes_flat, in_tree2 = tree_flatten(in_axes)\n",
    "    if in_tree != in_tree2: raise TypeError\n",
    "    f_flat, out_tree = flatten_fun(f, in_tree)\n",
    "    outs_flat = vmap_flat(f_flat, in_axes_flat, *args_flat)\n",
    "    return tree_unflatten(out_tree(), outs_flat)\n",
    "  return batched_f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_one_to_a_scalar(scalar):\n",
    "  assert np.ndim(scalar) == 0\n",
    "  return 1 + scalar\n",
    "\n",
    "vector_in = np.arange(3.)\n",
    "vector_out = vmap(add_one_to_a_scalar, (0,))(vector_in)\n",
    "\n",
    "print(vector_in)\n",
    "print(vector_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def jacfwd(f, x):\n",
    "  pushfwd = lambda v: jvp(f, (x,), (v,))[1]\n",
    "  vecs_in = np.eye(np.size(x)).reshape(np.shape(x) * 2)\n",
    "  return vmap(pushfwd, (0,))(vecs_in)\n",
    "\n",
    "def f(x):\n",
    "  return sin(x)\n",
    "\n",
    "jacfwd(f, np.arange(3.))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "That's it for `jvp` and `vmap`! Before moving on, let's highlight a few\n",
    "simplifications in what we've seen so far compared to the full JAX\n",
    "implementation:\n",
    "1. **Fewer, simpler primitives.** More primitives means more interpretation\n",
    "rules, and for more complex primitives (like for convolution or advanced\n",
    "indexing) each rule is harder to write. But the overarching design is no\n",
    "different.\n",
    "2. **No pytrees.** Transformations expect arrays in, and either a single array\n",
    "   out or a flat list of arrays out.\n",
    "3. **Missing optimization: no symbolic zeros in autodiff.**\n",
    "4. **No special call primitives yet.** The core machinery needs to be\n",
    "    generalized to handle the most flexible kind of higher-order primitive,\n",
    "    used by `jax.custom_jvp` and `jax.custom_vjp`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 2: Jaxprs\n",
    "\n",
    "The next transformations are the horizon are `jit` for just-in-time\n",
    "compilation and `vjp` for reverse-mode autodiff.  (`grad` is just a small\n",
    "wrapper around `vjp`.) Whereas `jvp` and `vmap` only needed each `Tracer` to\n",
    "carry a little bit of extra context, for both `jit` and `vjp` we need much\n",
    "richer context: we need to represent _programs_. That is, we need jaxprs!\n",
    "\n",
    "Jaxprs are JAX's internal intermediate representation of programs. They are\n",
    "explicitly typed, functional, first-order, and in ANF form. We need a\n",
    "program representation for `jit` because the purpose of `jit` is to stage\n",
    "computation out of Python. For any computation we want to stage out, we need\n",
    "to be able to represent it as data, and build it up as we trace a Python\n",
    "function. Similarly, `vjp` needs a way to represent the computation for the\n",
    "backward pass of reverse-mode autodiff. We use the same jaxpr program\n",
    "representation for both needs.\n",
    "\n",
    "(Building a program representation is the most\n",
    "[free](https://en.wikipedia.org/wiki/Free_object) kind of\n",
    "trace-transformation, and so except for issues around handling native Python\n",
    "control flow, any transformation could be implemented by first tracing to a\n",
    "jaxpr and then interpreting the jaxpr.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Jaxpr data strutures\n",
    "\n",
    "The jaxpr term syntax is roughly:\n",
    "\n",
    "```\n",
    "jaxpr ::=\n",
    "  { lambda <binder> , ... .\n",
    "    let <eqn>\n",
    "        ...\n",
    "    in ( <atom> , ... ) }\n",
    "\n",
    "binder ::= <var>:<array_type>\n",
    "var ::= a | b | c | ...\n",
    "atom ::= <var> | <literal>\n",
    "literal ::= <int32> | <float32>\n",
    "\n",
    "eqn ::= <binder> , ... = <primitive> [ <params> ] <atom> , ...\n",
    "```\n",
    "\n",
    "The syntax of types is:\n",
    "\n",
    "```\n",
    "jaxpr_type ::= [ <array_type> , ... ] -> [ <array_type> , ... ]\n",
    "array_type ::= <dtype>[<shape>]\n",
    "dtype ::= f32 | f64 | i32 | i64\n",
    "shape ::= <int> , ...\n",
    "```\n",
    "\n",
    "How do we represent these as Python data structures? We reuse ShapedArrays to\n",
    "represent types, and we can represent the term syntax with a few Python\n",
    "structs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_end_of_cell_marker": 0,
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "from typing import Set\n",
    "\n",
    "class Var:\n",
    "  aval: ShapedArray\n",
    "  def __init__(self, aval): self.aval = aval\n",
    "\n",
    "class Lit:\n",
    "  val: Any\n",
    "  aval: ShapedArray\n",
    "\n",
    "  def __init__(self, val):\n",
    "    self.val = val\n",
    "    self.aval = raise_to_shaped(get_aval(self.val))\n",
    "\n",
    "Atom = Union[Var, Lit]\n",
    "\n",
    "class JaxprEqn(NamedTuple):\n",
    "  primitive: Primitive\n",
    "  inputs: List[Atom]\n",
    "  params: Dict[str, Any]\n",
    "  out_binders: List[Var]\n",
    "\n",
    "class Jaxpr(NamedTuple):\n",
    "  in_binders: List[Var]\n",
    "  eqns: List[JaxprEqn]\n",
    "  outs: List[Atom]\n",
    "\n",
    "def raise_to_shaped(aval):\n",
    "  return ShapedArray(aval.shape, aval.dtype)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "Type-checking a jaxpr involves checking that there are no unbound variables,\n",
    "that variables are only bound once, and that for each equation the type of\n",
    "the primitive application matches the type of the output binders."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_end_of_cell_marker": 0,
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "class JaxprType:\n",
    "  in_types: List[ShapedArray]\n",
    "  out_type: List[ShapedArray]\n",
    "\n",
    "  def __init__(self, in_types, out_types):\n",
    "    self.in_types = in_types\n",
    "    self.out_types = out_types\n",
    "\n",
    "  def __repr__(self):\n",
    "    in_types = ', '.join(aval.str_short() for aval in self.in_types)\n",
    "    out_types = ', '.join(aval.str_short() for aval in self.out_types)\n",
    "    return f'({in_types}) -> ({out_types})'\n",
    "\n",
    "def typecheck_jaxpr(jaxpr: Jaxpr) -> JaxprType:\n",
    "  env: Set[Var] = set()\n",
    "\n",
    "  for v in jaxpr.in_binders:\n",
    "    if v in env: raise TypeError\n",
    "    env.add(v)\n",
    "\n",
    "  for eqn in jaxpr.eqns:\n",
    "    in_types = [typecheck_atom(env, x) for x in eqn.inputs]\n",
    "    out_types = abstract_eval_rules[eqn.primitive](*in_types, **eqn.params)\n",
    "    for out_binder, out_type in zip(eqn.out_binders, out_types):\n",
    "      if not types_equal(out_type, out_binder.aval): raise TypeError\n",
    "    for out_binder in eqn.out_binders:\n",
    "      if out_binder in env: raise TypeError\n",
    "      env.add(out_binder)\n",
    "\n",
    "  in_types = [v.aval for v in jaxpr.in_binders]\n",
    "  out_types = [typecheck_atom(env, x) for x in jaxpr.outs]\n",
    "  return JaxprType(in_types, out_types)\n",
    "\n",
    "def typecheck_atom(env: Set[Var], x: Atom) -> ShapedArray:\n",
    "  if isinstance(x, Var):\n",
    "    if x not in env: raise TypeError(\"unbound variable\")\n",
    "    return x.aval\n",
    "  elif isinstance(x, Lit):\n",
    "    return raise_to_shaped(get_aval(x.val))\n",
    "  else:\n",
    "    assert False\n",
    "\n",
    "def types_equal(a: ShapedArray, b: ShapedArray) -> bool:\n",
    "  return a.shape == b.shape and a.dtype == b.dtype"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can apply the function represented by a jaxpr to arguments with a simple\n",
    "interpreter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def eval_jaxpr(jaxpr: Jaxpr, args: List[Any]) -> List[Any]:\n",
    "  env: Dict[Var, Any] = {}\n",
    "\n",
    "  def read(x: Atom) -> Any:\n",
    "    return env[x] if type(x) is Var else x.val\n",
    "\n",
    "  def write(v: Var, val: Any) -> None:\n",
    "    env[v] = val\n",
    "\n",
    "  map(write, jaxpr.in_binders, args)\n",
    "  for eqn in jaxpr.eqns:\n",
    "    in_vals = map(read, eqn.inputs)\n",
    "    outs = bind(eqn.primitive, *in_vals, **eqn.params)\n",
    "    map(write, eqn.out_binders, outs)\n",
    "  return map(read, jaxpr.outs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def jaxpr_as_fun(jaxpr: Jaxpr):\n",
    "  return lambda *args: eval_jaxpr(jaxpr, args)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By using `bind` in the interpreter, this interpreter itself is traceable."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Building jaxprs with tracing\n",
    "\n",
    "Now that we have jaxprs as a data structure, we need ways to produce these\n",
    "from tracing Python code. In general there are two variants of how we trace to\n",
    "a jaxpr; `jit` uses one and `vjp` uses the other. We'll start with the one\n",
    "used by `jit`, which is also used by control flow primitives like `lax.cond`,\n",
    "`lax.while_loop`, and `lax.scan`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_end_of_cell_marker": 0,
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "# NB: the analogous class in JAX is called 'DynamicJaxprTracer'\n",
    "class JaxprTracer(Tracer):\n",
    "  __slots__ = ['aval']\n",
    "  aval: ShapedArray\n",
    "\n",
    "  def __init__(self, trace, aval):\n",
    "    self._trace = trace\n",
    "    self.aval = aval\n",
    "\n",
    "# NB: the analogous class in JAX is called 'DynamicJaxprTrace'\n",
    "class JaxprTrace(Trace):\n",
    "  def new_arg(self, aval: ShapedArray) -> JaxprTracer:\n",
    "    aval = raise_to_shaped(aval)\n",
    "    tracer = self.builder.new_tracer(self, aval)\n",
    "    self.builder.tracer_to_var[id(tracer)] = Var(aval)\n",
    "    return tracer\n",
    "\n",
    "  def get_or_make_const_tracer(self, val: Any) -> JaxprTracer:\n",
    "    tracer = self.builder.const_tracers.get(id(val))\n",
    "    if tracer is None:\n",
    "      tracer = self.builder.new_tracer(self, raise_to_shaped(get_aval(val)))\n",
    "      self.builder.add_const(tracer, val)\n",
    "    return tracer\n",
    "  pure = lift = get_or_make_const_tracer\n",
    "\n",
    "  def process_primitive(self, primitive, tracers, params):\n",
    "    avals_in = [t.aval for t in tracers]\n",
    "    avals_out = abstract_eval_rules[primitive](*avals_in, **params)\n",
    "    out_tracers = [self.builder.new_tracer(self, a) for a in avals_out]\n",
    "    inputs = [self.builder.getvar(t) for t in tracers]\n",
    "    outvars = [self.builder.add_var(t) for t in out_tracers]\n",
    "    self.builder.add_eqn(JaxprEqn(primitive, inputs, params, outvars))\n",
    "    return out_tracers\n",
    "\n",
    "  @property\n",
    "  def builder(self):\n",
    "    return self.main.global_data\n",
    "\n",
    "# NB: in JAX, instead of a dict we attach impl rules to the Primitive instance\n",
    "abstract_eval_rules = {}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that we keep as interpreter-global data a builder object, which keeps\n",
    "track of variables, constants, and eqns as we build up the jaxpr."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "class JaxprBuilder:\n",
    "  eqns: List[JaxprEqn]\n",
    "  tracer_to_var: Dict[int, Var]\n",
    "  const_tracers: Dict[int, JaxprTracer]\n",
    "  constvals: Dict[Var, Any]\n",
    "  tracers: List[JaxprTracer]\n",
    "\n",
    "  def __init__(self):\n",
    "    self.eqns = []\n",
    "    self.tracer_to_var = {}\n",
    "    self.const_tracers = {}\n",
    "    self.constvals = {}\n",
    "    self.tracers = []\n",
    "\n",
    "  def new_tracer(self, trace: JaxprTrace, aval: ShapedArray) -> JaxprTracer:\n",
    "    tracer = JaxprTracer(trace, aval)\n",
    "    self.tracers.append(tracer)\n",
    "    return tracer\n",
    "\n",
    "  def add_eqn(self, eqn: JaxprEqn) -> None:\n",
    "    self.eqns.append(eqn)\n",
    "\n",
    "  def add_var(self, tracer: JaxprTracer) -> Var:\n",
    "    assert id(tracer) not in self.tracer_to_var\n",
    "    var = self.tracer_to_var[id(tracer)] = Var(tracer.aval)\n",
    "    return var\n",
    "\n",
    "  def getvar(self, tracer: JaxprTracer) -> Var:\n",
    "    var = self.tracer_to_var.get(id(tracer))\n",
    "    assert var is not None\n",
    "    return var\n",
    "\n",
    "  def add_const(self, tracer: JaxprTracer, val: Any) -> Var:\n",
    "    var = self.add_var(tracer)\n",
    "    self.const_tracers[id(val)] = tracer\n",
    "    self.constvals[var] = val\n",
    "    return var\n",
    "\n",
    "  def build(self, in_tracers: List[JaxprTracer], out_tracers: List[JaxprTracer]\n",
    "            ) -> Tuple[Jaxpr, List[Any]]:\n",
    "    constvars, constvals = unzip2(self.constvals.items())\n",
    "    t2v = lambda t: self.tracer_to_var[id(t)]\n",
    "    in_binders = constvars + [t2v(t) for t in in_tracers]\n",
    "    out_vars = [t2v(t) for t in out_tracers]\n",
    "    jaxpr = Jaxpr(in_binders, self.eqns, out_vars)\n",
    "    typecheck_jaxpr(jaxpr)\n",
    "    return jaxpr, constvals"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The rules we need for `JaxprTrace.process_primitive` are essentially typing\n",
    "rules for primitive applications: given the primitive, its parameters, and\n",
    "types for the inputs, the rule must produce a type for the output, which is\n",
    "then packaged with the output `JaxprTracer`. We can use abstract evaluation\n",
    "rules for this same purpose, even though they can be more general (since\n",
    "abstract evaluation rules must accept ConcreteArray inputs, and since they\n",
    "need only return an upper bound on the set of possible outputs, they can\n",
    "produce ConcreteArray outputs as well). We'll reuse these abstract evaluation\n",
    "rules for the other jaxpr-producing trace machinery, where the potential extra\n",
    "generality is useful."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def broadcast_shapes(*shapes):\n",
    "  assert len(shapes) > 1\n",
    "  for sizes in zip(*shapes):\n",
    "    sizes = [d for d in sizes if d != 1]\n",
    "    if sizes[:-1] != sizes[1:]:\n",
    "      raise Exception\n",
    "  return tuple(next((d for d in sizes if d != 1), 1) for sizes in zip(*shapes))\n",
    "\n",
    "def broadcasting_binop_abstract_eval_rule(*avals_in):\n",
    "  out_dtype = np.result_type(*map(np.result_type, avals_in))\n",
    "  out_shape = broadcast_shapes(*map(np.shape, avals_in))\n",
    "  return [ShapedArray(out_shape, out_dtype)]\n",
    "\n",
    "abstract_eval_rules[add_p] = broadcasting_binop_abstract_eval_rule\n",
    "abstract_eval_rules[mul_p] = broadcasting_binop_abstract_eval_rule\n",
    "\n",
    "def vectorized_unop_abstract_eval_rule(aval_in):\n",
    "  return [ShapedArray(np.shape(aval_in), np.result_type(aval_in))]\n",
    "\n",
    "abstract_eval_rules[sin_p] = vectorized_unop_abstract_eval_rule\n",
    "abstract_eval_rules[cos_p] = vectorized_unop_abstract_eval_rule\n",
    "abstract_eval_rules[neg_p] = vectorized_unop_abstract_eval_rule\n",
    "\n",
    "def reduce_sum_abstract_eval_rule(aval_in, *, axis):\n",
    "  new_shape = [d for i, d in enumerate(aval_in.shape) if i != axis]\n",
    "  return [ShapedArray(tuple(new_shape), aval_in.dtype)]\n",
    "abstract_eval_rules[reduce_sum_p] = reduce_sum_abstract_eval_rule\n",
    "\n",
    "def broadcast_abstract_eval(x, *, shape, axes):\n",
    "  return [ShapedArray(tuple(shape), np.result_type(x))]\n",
    "abstract_eval_rules[broadcast_p] = broadcast_abstract_eval"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To check our implementation of jaxprs, we can add a `make_jaxpr`\n",
    "transformation and a pretty-printer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "from functools import lru_cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "@lru_cache()\n",
    "def make_jaxpr_v1(f, *avals_in):\n",
    "  avals_in, in_tree = tree_flatten(avals_in)\n",
    "  f, out_tree = flatten_fun(f, in_tree)\n",
    "\n",
    "  builder = JaxprBuilder()\n",
    "  with new_main(JaxprTrace, builder) as main:\n",
    "    trace = JaxprTrace(main)\n",
    "    tracers_in = [trace.new_arg(aval) for aval in avals_in]\n",
    "    outs = f(*tracers_in)\n",
    "    tracers_out = [full_raise(trace, out) for out in outs]\n",
    "    jaxpr, consts = builder.build(tracers_in, tracers_out)\n",
    "  return jaxpr, consts, out_tree()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_end_of_cell_marker": 0,
    "lines_to_next_cell": 1,
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "import string\n",
    "\n",
    "class PPrint:\n",
    "  lines: List[Tuple[int, str]]\n",
    "\n",
    "  def __init__(self, lines):\n",
    "    self.lines = lines\n",
    "\n",
    "  def indent(self, indent: int) -> 'PPrint':\n",
    "    return PPrint([(indent + orig_indent, s) for orig_indent, s in self.lines])\n",
    "\n",
    "  def __add__(self, rhs: 'PPrint') -> 'PPrint':\n",
    "    return PPrint(self.lines + rhs.lines)\n",
    "\n",
    "  def __rshift__(self, rhs: 'PPrint') -> 'PPrint':\n",
    "    if not rhs.lines: return self\n",
    "    if not self.lines: return rhs\n",
    "    indent, s = self.lines[-1]\n",
    "    indented_block = rhs.indent(indent + len(s))\n",
    "    common_line = s + ' ' * rhs.lines[0][0] + rhs.lines[0][1]\n",
    "    return PPrint(self.lines[:-1]\n",
    "                  + [(indent, common_line)]\n",
    "                  + indented_block.lines[1:])\n",
    "\n",
    "  def __str__(self) -> str:\n",
    "    return '\\n'.join(' ' * indent + s for indent, s in self.lines)\n",
    "\n",
    "def pp(s: Any) -> PPrint:\n",
    "  return PPrint([(0, line) for line in str(s).splitlines()])\n",
    "\n",
    "def vcat(ps: List[PPrint]) -> PPrint:\n",
    "  return sum(ps, pp(''))\n",
    "\n",
    "def pp_jaxpr(jaxpr: Jaxpr):\n",
    "  namegen = (''.join(s) for r in it.count(1)\n",
    "             for s in it.permutations(string.ascii_lowercase, r))\n",
    "  names = defaultdict(lambda: next(namegen))\n",
    "  in_binders = ', '.join(var_str(names, x) for x in jaxpr.in_binders)\n",
    "  eqns = vcat([pp_eqn(names, e) for e in jaxpr.eqns])\n",
    "  outs = ', '.join(names[v] if isinstance(v, Var) else str(v.val)\n",
    "                   for v in jaxpr.outs)\n",
    "  return (pp(f'{{ lambda {in_binders} .') +\n",
    "          ((pp('let ') >> eqns) + pp(f'in ( {outs} ) }}')).indent(2))\n",
    "\n",
    "def var_str(names: Dict[Var, str], v: Var) -> str:\n",
    "  return f'{names[v]}:{v.aval.str_short()}'\n",
    "\n",
    "def pp_eqn(names: Dict[Var, str], eqn: JaxprEqn) -> PPrint:\n",
    "  lhs = pp(' '.join(var_str(names, v) for v in eqn.out_binders))\n",
    "  rhs = (pp(eqn.primitive.name) >> pp_params(eqn.params) >>\n",
    "         pp(' '.join(names[x] if isinstance(x, Var) else str(x.val)\n",
    "                     for x in eqn.inputs)))\n",
    "  return lhs >> pp(' = ') >> rhs\n",
    "\n",
    "def pp_params(params: Dict[str, Any]) -> PPrint:\n",
    "  items = sorted(params.items())\n",
    "  if items:\n",
    "    return pp(' [ ') >> vcat([pp(f'{k}={v}') for k, v in items]) >> pp(' ] ')\n",
    "  else:\n",
    "    return pp(' ')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "jaxpr, consts, _ = make_jaxpr_v1(lambda x: 2. * x, raise_to_shaped(get_aval(3.)))\n",
    "print(pp_jaxpr(jaxpr))\n",
    "print(typecheck_jaxpr(jaxpr))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But there's a limitation here: because of how `find_top_trace` operates by\n",
    "data dependence, `make_jaxpr_v1` can't stage out all the primitive operations\n",
    "performed by the Python callable it's given. For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "jaxpr, consts, _ = make_jaxpr_v1(lambda: mul(2., 2.))\n",
    "print(pp_jaxpr(jaxpr))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is precisely the issue that\n",
    "[omnistaging](https://github.com/google/jax/pull/3370) fixed.\n",
    "We want to ensure that the `JaxprTrace` started by `make_jaxpr` is always\n",
    "applied, regardless of whether any inputs to `bind` are boxed in corresponding\n",
    "`JaxprTracer` instances. We can achieve this by employing the `dynamic_trace`\n",
    "global defined in Part 1:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "@contextmanager\n",
    "def new_dynamic(main: MainTrace):\n",
    "  global dynamic_trace\n",
    "  prev_dynamic_trace, dynamic_trace = dynamic_trace, main\n",
    "  try:\n",
    "    yield\n",
    "  finally:\n",
    "    dynamic_trace = prev_dynamic_trace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "@lru_cache()  # ShapedArrays are hashable\n",
    "def make_jaxpr(f, *avals_in):\n",
    "  avals_in, in_tree = tree_flatten(avals_in)\n",
    "  f, out_tree = flatten_fun(f, in_tree)\n",
    "\n",
    "  builder = JaxprBuilder()\n",
    "  with new_main(JaxprTrace, builder) as main:\n",
    "    with new_dynamic(main):\n",
    "      trace = JaxprTrace(main)\n",
    "      tracers_in = [trace.new_arg(aval) for aval in avals_in]\n",
    "      outs = f(*tracers_in)\n",
    "      tracers_out = [full_raise(trace, out) for out in outs]\n",
    "      jaxpr, consts = builder.build(tracers_in, tracers_out)\n",
    "  return jaxpr, consts, out_tree()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_end_of_cell_marker": 0,
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "jaxpr, consts, _ = make_jaxpr(lambda: mul(2., 2.))\n",
    "print(pp_jaxpr(jaxpr))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using `dynamic_trace` this way is conceptually the same as stashing the\n",
    "current interpreter stack and starting a new one with the `JaxprTrace` at the\n",
    "bottom. That is, no interpreters lower in the stack than the `dynamic_trace`\n",
    "are applied (since `JaxprTrace.process_primitive` doesn't call `bind`), though\n",
    "if the Python callable being traced to a jaxpr itself uses transformations\n",
    "then those can be pushed onto the interpreter stack above the `JaxprTrace`.\n",
    "But temporarily stashing the interpreter stack would break up the system\n",
    "state. The `dynamic_trace` tag achieves the same goals while keeping the\n",
    "system state simpler."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's it for jaxprs! With jaxprs in hand, we can implement the remaining\n",
    "major JAX features. But before moving on, let's highlight some\n",
    "simplifications we've made:\n",
    "1. **Single-output primitives and jaxprs.**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 3: `jit`, simplified\n",
    "\n",
    "While `jit` has a transformation-like API in that it accepts a Python callable\n",
    "as an argument, under the hood it's really a higher-order primitive rather\n",
    "than a transformation. A primitive is _higher-order_ when it's parameterized\n",
    "by a function."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### \"Final style\" and \"initial style\"\n",
    "\n",
    "There are two options for how to handle higher-order primitives. Each requires\n",
    "a different approach to tracing and engenders different tradeoffs:\n",
    "1. **`bind` takes a Python callable as an argument.** We defer forming a jaxpr\n",
    "   until as late as possible, namely until we're running the final interpreter\n",
    "   at the bottom of the interpreter stack. That way we can swap a `JaxprTrace`\n",
    "   in at the bottom of the interpreter stack and thus stage out rather than\n",
    "   execute all primitive operations. With this approach, transformations in\n",
    "   the stack get applied as we execute the Python callable as usual. This\n",
    "   approach can be very tricky to implement, but it's as general as possible\n",
    "   because it allows higher-order primitives not to raise the abstraction\n",
    "   level of their arguments and thus allows data-dependent Python control\n",
    "   flow. We refer to this approach as using a \"final-style higher-order\n",
    "   primitive\" employing the discharge-at-tracing-time \"final-style\n",
    "   transformations\" we've used so far.\n",
    "2. **`bind` takes a jaxpr as an argument.** Before we call `bind`, in the\n",
    "   primitive wrapper we can just use `make_jaxpr` to form a jaxpr up-front and\n",
    "   be done with the Python callable entirely. In this case, `make_jaxpr` puts\n",
    "   its `JaxprTrace` at the top of the interpreter stack, and no\n",
    "   transformations lower in the stack, which might enter via closed-over\n",
    "   Tracers, are applied to the Python callable as we trace it.\n",
    "   (Transformations applied within the Python callable are applied as usual,\n",
    "   being added to the stack above the JaxprTrace.) Instead, the\n",
    "   transformations lower in the stack are later applied to the call primitive,\n",
    "   and the call primitive's rules must then transform the jaxpr itself.\n",
    "   Because we trace to a jaxpr up-front, this approach can't support\n",
    "   data-dependent Python control flow, but it is more straightforward to\n",
    "   implement. We refer to this kind of higher-order primitive as an\n",
    "   \"initial-style higher-order primitive\", and say that its jaxpr-processing\n",
    "   transformation rules are \"initial-style transformation rules.\"\n",
    "\n",
    "The latter approach fits for `jit` because we don't need to support\n",
    "data-dependent Python control flow in the user-provided Python callable, as\n",
    "the whole purpose of `jit` is to stage computation out of Python to be\n",
    "executed by XLA. (In contrast, `custom_jvp` is a higher-order primitive in\n",
    "which we want to support data-dependent Python control flow.)\n",
    "\n",
    "Historically, we started using the \"initial-style\" and \"final-style\"\n",
    "terminology after reading the [typed tagless final\n",
    "interpreters](http://okmij.org/ftp/tagless-final/index.html) paper, and\n",
    "jokingly referring to JAX as an implementation of \"untyped tagful final\n",
    "interpreters.\" We don't claim to carry over (or understand) any deep meaning\n",
    "behind these terms; we loosely use \"initial style\" to mean \"build an AST and\n",
    "then transform it\", and we use \"final style\" to mean \"transform as we trace.\"\n",
    "But it's just imprecise yet sticky jargon."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With the initial-style approach, here's the user-facing `jit` wrapper:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_end_of_cell_marker": 0,
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def jit(f):\n",
    "  def f_jitted(*args):\n",
    "    avals_in = [raise_to_shaped(get_aval(x)) for x in args]\n",
    "    jaxpr, consts, out_tree = make_jaxpr(f, *avals_in)\n",
    "    outs = bind(xla_call_p, *consts, *args, jaxpr=jaxpr, num_consts=len(consts))\n",
    "    return tree_unflatten(out_tree, outs)\n",
    "  return f_jitted\n",
    "\n",
    "xla_call_p = Primitive('xla_call')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With any new primitive, we need to give it transformation rules, starting with\n",
    "its evaluation rule. When we evaluate an application of the `xla_call`\n",
    "primitive, we want to stage out out the computation to XLA. That involves\n",
    "translating the jaxpr to an XLA HLO program, transferring the argument values\n",
    "to the XLA device, executing the XLA program, and transferring back the\n",
    "results. We'll cache the XLA HLO compilation so that for each `jit`ted\n",
    "function it only needs to be performed once per argument shape and dtype\n",
    "signature.\n",
    "\n",
    "First, some utilities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "class IDHashable:\n",
    "  val: Any\n",
    "\n",
    "  def __init__(self, val):\n",
    "    self.val = val\n",
    "\n",
    "  def __hash__(self) -> int:\n",
    "    return id(self.val)\n",
    "\n",
    "  def __eq__(self, other):\n",
    "    return type(other) is IDHashable and id(self.val) == id(other.val)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we'll define the evaluation rule for `xla_call`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_end_of_cell_marker": 0,
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "from jax.lib import xla_bridge as xb\n",
    "from jax.lib import xla_client as xc\n",
    "xe = xc._xla\n",
    "xops = xc._xla.ops\n",
    "\n",
    "def xla_call_impl(*args, jaxpr: Jaxpr, num_consts: int):\n",
    "  consts, args = args[:num_consts], args[num_consts:]\n",
    "  hashable_consts = tuple(map(IDHashable, consts))\n",
    "  execute = xla_callable(IDHashable(jaxpr), hashable_consts)\n",
    "  return execute(*args)\n",
    "impl_rules[xla_call_p] = xla_call_impl\n",
    "\n",
    "@lru_cache()\n",
    "def xla_callable(hashable_jaxpr: IDHashable, hashable_consts: Tuple[IDHashable]):\n",
    "  jaxpr: Jaxpr = hashable_jaxpr.val\n",
    "  consts = [x.val for x in hashable_consts]\n",
    "  in_avals = [v.aval for v in jaxpr.in_binders[len(consts):]]\n",
    "  c = xb.make_computation_builder('xla_call')\n",
    "  xla_consts = _xla_consts(c, consts)\n",
    "  xla_params = _xla_params(c, in_avals)\n",
    "  outs = jaxpr_subcomp(c, jaxpr, xla_consts + xla_params)\n",
    "  out = xops.Tuple(c, outs)\n",
    "  compiled = xb.get_backend(None).compile(c.build(out))\n",
    "  return partial(execute_compiled, compiled, [v.aval for v in jaxpr.outs])\n",
    "\n",
    "def _xla_consts(c: xe.XlaBuilder, consts: List[Any]) -> List[xe.XlaOp]:\n",
    "  unique_consts = {id(cnst): cnst for cnst in consts}\n",
    "  xla_consts = {\n",
    "      id_: xops.ConstantLiteral(c, cnst) for id_, cnst in unique_consts.items()}\n",
    "  return [xla_consts[id(cnst)] for cnst in consts]\n",
    "\n",
    "def _xla_params(c: xe.XlaBuilder, avals_in: List[ShapedArray]) -> List[xe.XlaOp]:\n",
    "  return [xb.parameter(c, i, _xla_shape(a)) for i, a in enumerate(avals_in)]\n",
    "\n",
    "def _xla_shape(aval: ShapedArray) -> xe.Shape:\n",
    "  return xc.Shape.array_shape(xc.dtype_to_etype(aval.dtype), aval.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The main action is in `xla_callable`, which compiles a jaxpr into an XLA HLO\n",
    "program using `jaxpr_subcomp`, then returns a callable which executes the\n",
    "compiled program:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_end_of_cell_marker": 0,
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def jaxpr_subcomp(c: xe.XlaBuilder, jaxpr: Jaxpr, args: List[xe.XlaOp]\n",
    "                  ) -> xe.XlaOp:\n",
    "  env: Dict[Var, xe.XlaOp] = {}\n",
    "\n",
    "  def read(x: Atom) -> xe.XlaOp:\n",
    "    return env[x] if type(x) is Var else xb.constant(c, x.val)\n",
    "\n",
    "  def write(v: Var, val: xe.XlaOp) -> None:\n",
    "    env[v] = val\n",
    "\n",
    "  map(write, jaxpr.in_binders, args)\n",
    "  for eqn in jaxpr.eqns:\n",
    "    in_avals = [x.aval for x in eqn.inputs]\n",
    "    in_vals = map(read, eqn.inputs)\n",
    "    rule = xla_translations[eqn.primitive]\n",
    "    out_vals = rule(c, in_avals, in_vals, **eqn.params)\n",
    "    map(write, eqn.out_binders, out_vals)\n",
    "  return map(read, jaxpr.outs)\n",
    "\n",
    "def execute_compiled(compiled, out_avals, *args):\n",
    "  input_bufs = [input_handlers[type(x)](x) for x in args]\n",
    "  out_bufs = compiled.execute(input_bufs)\n",
    "  return [handle_result(aval, buf) for aval, buf in zip(out_avals, out_bufs)]\n",
    "\n",
    "input_handlers = {\n",
    "    int: xb.get_backend(None).buffer_from_pyval,\n",
    "    float: xb.get_backend(None).buffer_from_pyval,\n",
    "    np.ndarray: xb.get_backend(None).buffer_from_pyval,\n",
    "}\n",
    "\n",
    "def handle_result(aval: ShapedArray, buf):\n",
    "  del aval  # Unused for now.\n",
    "  return buf.to_py()\n",
    "\n",
    "xla_translations = {}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that `jaxpr_subcomp` has the structure of a simple interpreter. That's\n",
    "a common pattern: the way we process jaxprs is usually with an interpreter.\n",
    "And as with any interpreter, we need an interpretation rule for each\n",
    "primitive:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def direct_translation(op, c, in_avals, in_vals):\n",
    "  del c, in_avals\n",
    "  return [op(*in_vals)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "xla_translations[add_p] = partial(direct_translation, xops.Add)\n",
    "xla_translations[mul_p] = partial(direct_translation, xops.Mul)\n",
    "xla_translations[neg_p] = partial(direct_translation, xops.Neg)\n",
    "xla_translations[sin_p] = partial(direct_translation, xops.Sin)\n",
    "xla_translations[cos_p] = partial(direct_translation, xops.Cos)\n",
    "xla_translations[greater_p] = partial(direct_translation, xops.Gt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def reduce_sum_translation(c, in_avals, in_vals, *, axis):\n",
    "  (x_aval,), (x,) = in_avals, in_vals\n",
    "  zero = xops.ConstantLiteral(c, np.array(0, x_aval.dtype))\n",
    "  subc = xb.make_computation_builder('add')\n",
    "  shape = _xla_shape(ShapedArray((), x_aval.dtype))\n",
    "  xops.Add(xops.Parameter(subc, 0, shape), xops.Parameter(subc, 1, shape))\n",
    "  return [xops.Reduce(c, [x], [zero], subc.build(), [axis])]\n",
    "xla_translations[reduce_sum_p] = reduce_sum_translation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def broadcast_translation(c, in_avals, in_vals, *, shape, axes):\n",
    "  x, = in_vals\n",
    "  dims_complement = [i for i in range(len(shape)) if i not in axes]\n",
    "  return [xops.BroadcastInDim(x, shape, dims_complement)]\n",
    "xla_translations[broadcast_p] = broadcast_translation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With that, we can now use `jit` to stage out, compile, and execute programs\n",
    "with XLA!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "@jit\n",
    "def f(x, y):\n",
    "  print('tracing!')\n",
    "  return sin(x) * cos(y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z = f(3., 4.)  # 'tracing!' prints the first time\n",
    "print(z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "z = f(4., 5.)  # 'tracing!' doesn't print, compilation cache hit!\n",
    "print(z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "@jit\n",
    "def f(x):\n",
    "  return reduce_sum(x, axis=0)\n",
    "\n",
    "print(f(np.array([1., 2., 3.])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_end_of_cell_marker": 0,
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def f(x):\n",
    "  y = sin(x) * 2.\n",
    "  z = - y + x\n",
    "  return z\n",
    "\n",
    "def deriv(f):\n",
    "  return lambda x: jvp(f, (x,), (1.,))[1]\n",
    "\n",
    "print(    deriv(deriv(f))(3.))\n",
    "print(jit(deriv(deriv(f)))(3.))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Instead of implementing `jit` to first trace to a jaxpr and then to lower the\n",
    "jaxpr to XLA HLO, it might appear that we could have skipped the jaxpr step\n",
    "and just lowered to HLO while tracing. That is, perhaps we could have instead\n",
    "implemented `jit` with a `Trace` and `Tracer` that appended to the XLA HLO\n",
    "graph incrementally on each primitive bind. That's correct for now, but won't\n",
    "be possible when we introduce compiled SPMD computations because there we must\n",
    "know the number of replicas needed before compiling the program."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We haven't yet defined any transformation rules for `xla_call_p` other than\n",
    "its evaluation rule. That is, we can't yet do `vmap`-of-`jit` or\n",
    "`jvp`-of-`jit` or even `jit`-of`-jit`. Instead `jit` has to be at the \"top\n",
    "level.\" Let's fix that!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def xla_call_jvp_rule(primals, tangents, *, jaxpr, num_consts):\n",
    "  del num_consts  # Unused.\n",
    "  new_jaxpr, new_consts = jvp_jaxpr(jaxpr)\n",
    "  outs = bind(xla_call_p, *new_consts, *primals, *tangents, jaxpr=new_jaxpr,\n",
    "              num_consts=len(new_consts))\n",
    "  n = len(outs) // 2\n",
    "  primals_out, tangents_out = outs[:n], outs[n:]\n",
    "  return primals_out, tangents_out\n",
    "jvp_rules[xla_call_p] = xla_call_jvp_rule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def jvp_jaxpr(jaxpr: Jaxpr) -> Tuple[Jaxpr, List[Any]]:\n",
    "  def jvp_traceable(*primals_and_tangents):\n",
    "    n = len(primals_and_tangents) // 2\n",
    "    primals, tangents = primals_and_tangents[:n], primals_and_tangents[n:]\n",
    "    return jvp(jaxpr_as_fun(jaxpr), primals, tangents)\n",
    "\n",
    "  in_avals = [v.aval for v in jaxpr.in_binders]\n",
    "  new_jaxpr, new_consts, _ = make_jaxpr(jvp_traceable, *in_avals, *in_avals)\n",
    "  return new_jaxpr, new_consts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "def xla_call_vmap_rule(axis_size, vals_in, dims_in, *, jaxpr, num_consts):\n",
    "  del num_consts  # Unused.\n",
    "  new_jaxpr, new_consts = vmap_jaxpr(jaxpr, axis_size, dims_in)\n",
    "  outs = bind(xla_call_p, *new_consts, *vals_in, jaxpr=new_jaxpr,\n",
    "              num_consts=len(new_consts))\n",
    "  return outs, [0] * len(outs)\n",
    "vmap_rules[xla_call_p] = xla_call_vmap_rule\n",
    "\n",
    "def vmap_jaxpr(jaxpr: Jaxpr, axis_size: int, bdims_in: List[BatchAxis]\n",
    "               ) -> Tuple[Jaxpr, List[Any]]:\n",
    "  vmap_traceable = vmap(jaxpr_as_fun(jaxpr), tuple(bdims_in))\n",
    "  in_avals = [unmapped_aval(axis_size, d, v.aval)\n",
    "              for v, d in zip(jaxpr.in_binders, bdims_in)]\n",
    "  new_jaxpr, new_consts, _ = make_jaxpr(vmap_traceable, *in_avals)\n",
    "  return new_jaxpr, new_consts\n",
    "\n",
    "def unmapped_aval(axis_size: int, batch_dim: BatchAxis, aval: ShapedArray\n",
    "                  ) -> ShapedArray:\n",
    "  if batch_dim is not_mapped:\n",
    "    return aval\n",
    "  else:\n",
    "    shape = list(aval.shape)\n",
    "    shape.insert(batch_dim, axis_size)\n",
    "    return ShapedArray(tuple(shape), aval.dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_end_of_cell_marker": 0,
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "@jit\n",
    "def f(x):\n",
    "  y = sin(x) * 2.\n",
    "  z = - y + x\n",
    "  return z\n",
    "\n",
    "x, xdot = 3., 1.\n",
    "y, ydot = jvp(f, (x,), (xdot,))\n",
    "print(y)\n",
    "print(ydot)\n",
    "\n",
    "ys = vmap(f, (0,))(np.arange(3.))\n",
    "print(ys)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One piece missing is device memory persistence for arrays. That is, we've\n",
    "defined `handle_result` to transfer results back to CPU memory as NumPy\n",
    "arrays, but it's often preferrable to avoid transferring results just to\n",
    "transfer them back for the next operation. We can do that by introducing a\n",
    "`DeviceArray` class, which can wrap XLA buffers and otherwise duck-type\n",
    "`numpy.ndarray`s:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def handle_result(aval: ShapedArray, buf):  # noqa: F811\n",
    "  return DeviceArray(aval, buf)\n",
    "\n",
    "class DeviceArray:\n",
    "  buf: Any\n",
    "  aval: ShapedArray\n",
    "\n",
    "  def __init__(self, aval, buf):\n",
    "    self.aval = aval\n",
    "    self.buf = buf\n",
    "\n",
    "  dtype = property(lambda self: self.aval.dtype)\n",
    "  shape = property(lambda self: self.aval.shape)\n",
    "  ndim  = property(lambda self: self.aval.ndim)\n",
    "\n",
    "  def __array__(self): return self.buf.to_py()\n",
    "  def __repr__(self):  return repr(self.buf.to_py())\n",
    "  def __str__(self):   return str(self.buf.to_py())\n",
    "\n",
    "  _neg = staticmethod(neg)\n",
    "  _add = staticmethod(add)\n",
    "  _radd = staticmethod(add)\n",
    "  _mul = staticmethod(mul)\n",
    "  _rmul = staticmethod(mul)\n",
    "  _gt = staticmethod(greater)\n",
    "input_handlers[DeviceArray] = lambda x: x.buf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit\n",
    "def f(x):\n",
    "  y = sin(x) * 2.\n",
    "  z = - y + x\n",
    "  return z\n",
    "\n",
    "x, xdot = 3., 1.\n",
    "y, ydot = jvp(f, (x,), (xdot,))\n",
    "print(y)\n",
    "print(ydot)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "last_runtime": {
    "build_target": "//learning/deepmind/dm_python:dm_notebook3",
    "kind": "private"
   },
   "name": "Autodidax: JAX core from scratch",
   "provenance": [],
   "toc_visible": true
  },
  "jupytext": {
   "formats": "ipynb,md:myst,py",
   "main_language": "python"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}