{ "cells": [ { "cell_type": "markdown", "id": "55dedd56", "metadata": {}, "source": [ "# Derivatives\n", "\n", "Basic idea: `jvp` and `vjp` as pushforward and pullback on $\\mathbb R^n$." ] }, { "cell_type": "code", "execution_count": 1, "id": "5da78d72", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:52.349252Z", "iopub.status.busy": "2026-07-26T13:35:52.348905Z", "iopub.status.idle": "2026-07-26T13:35:52.687481Z", "shell.execute_reply": "2026-07-26T13:35:52.687018Z" } }, "outputs": [], "source": [ "import jax\n", "import jax.numpy as jnp\n", "\n", "key = jax.random.key(42)" ] }, { "cell_type": "markdown", "id": "0dafe7f4", "metadata": {}, "source": [ "Example: $f: \\mathbb R^n \\to \\mathbb R^2$.\n", "\n", "[Note: python objects do have types but functions don't declare or check them. Python does let us annotate functions, but the type would be `jax.Array` without specifying $n$ here. There *are* [packages](https://github.com/patrick-kidger/jaxtyping) for annotating the latter.]" ] }, { "cell_type": "code", "execution_count": 2, "id": "ba0848e0", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:52.688853Z", "iopub.status.busy": "2026-07-26T13:35:52.688755Z", "iopub.status.idle": "2026-07-26T13:35:52.690474Z", "shell.execute_reply": "2026-07-26T13:35:52.690106Z" } }, "outputs": [], "source": [ "f = lambda x: jnp.array([\n", " jnp.sin(jnp.sum(x)),\n", " jnp.sum(x)]\n", ")" ] }, { "cell_type": "code", "execution_count": 3, "id": "d3d83e3b", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:52.691380Z", "iopub.status.busy": "2026-07-26T13:35:52.691315Z", "iopub.status.idle": "2026-07-26T13:35:52.760115Z", "shell.execute_reply": "2026-07-26T13:35:52.759640Z" } }, "outputs": [ { "data": { "text/plain": [ "Array([0.67023855, 0.7345302 ], dtype=float32)" ] }, "execution_count": 3, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# let's make it n = 3\n", "x = jax.random.normal(key, (3,))\n", "\n", "f(x)" ] }, { "cell_type": "markdown", "id": "e4f2470b", "metadata": {}, "source": [ "Everything's quite trivial here, but let's call it $f: M \\to N$ so we can write the input tangent space as $TM$ (instead of immediately identifying $T\\mathbb R^n \\cong \\mathbb R^n$).\n", "\n", "### Forward ~ Pushforward\n", "\n", "Then \"forward mode\" autodiff is\n", "$$\n", "\\begin{align}\n", "\\mathrm{jvp}: (M \\to N) \\times M \\times TM &\\to N \\times TN \\\\\n", " (f, x, v) &\\mapsto (f(x), df|_x(v)) \\,,\n", "\\end{align}\n", "$$\n", "i.e. the \"Jacobian * vector\" product = jvp.\n", "Note that we specify all three objects simultaneously (the function, the input, and the tangent). That is because in forward mode all are computed alongside each other (autodiff = automated chain rule; in forward mode input and chain rule go in the same direction)." ] }, { "cell_type": "code", "execution_count": 4, "id": "0f18e536", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:52.761231Z", "iopub.status.busy": "2026-07-26T13:35:52.761168Z", "iopub.status.idle": "2026-07-26T13:35:52.792384Z", "shell.execute_reply": "2026-07-26T13:35:52.791940Z" } }, "outputs": [ { "data": { "text/plain": [ "(Array([0.67023855, 0.7345302 ], dtype=float32),\n", " Array([0.7421457, 1. ], dtype=float32))" ] }, "execution_count": 4, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# note that this must be floats; [1, 0, 0] would give integer and jax would complain\n", "v = jnp.array([1.0, 0.0, 0.0])\n", "\n", "jax.jvp(f, (x,), (v,))" ] }, { "cell_type": "markdown", "id": "95734595", "metadata": {}, "source": [ "Note $x$ and $v$ are wrapped in a tuple because that's how jax deals with the issue that\n", "functions usually take multiple arguments.\n", "So for some other $g(x, y, z)$ the input space is $M = X \\times Y \\times Z$.\n", "For a single input, we still think of it as a \"tuple\" of one element.\n", "\n", "### Reverse ~ Pullback\n", "\n", "For reverse mode, we instead propagate the cotangent through the chain rule, which goes the opposite direction from the input.\n", "Thus, `vjp` only needs us to give $f$ and $x$, and then returns a *function* to compute the pullback later.\n", "$$\n", "\\begin{align}\n", "\\mathrm{vjp}: (M \\to N) \\times M &\\to N \\times (T^*N \\to T^*M) \\\\\n", " (f, x) &\\mapsto (f(x), (w \\mapsto w \\circ df|_{x}))\n", "\\end{align}\n", "$$\n", "Note the direction: `vjp` differentiates a map $M \\to N$, but the function it hands back goes the *other* way, from cotangents on $N$ to cotangents on $M$. That is what makes it a pullback." ] }, { "cell_type": "code", "execution_count": 5, "id": "3fc39ad6", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:52.793637Z", "iopub.status.busy": "2026-07-26T13:35:52.793551Z", "iopub.status.idle": "2026-07-26T13:35:52.819569Z", "shell.execute_reply": "2026-07-26T13:35:52.819164Z" } }, "outputs": [ { "data": { "text/plain": [ "Array([0.67023855, 0.7345302 ], dtype=float32)" ] }, "execution_count": 5, "metadata": {}, "output_type": "execute_result" } ], "source": [ "w = jnp.array([1.0, 1.0])\n", "\n", "# note that we do *not* wrap (x,) here.\n", "# by convention jax expects us to write it out, so it would be vjp(g, x, y, z) if\n", "# we had a function with multiple input arguments\n", "f_x, df_back = jax.vjp(f, x)\n", "f_x # consistency check; same output f(x) again" ] }, { "cell_type": "code", "execution_count": 6, "id": "635732b1", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:52.820814Z", "iopub.status.busy": "2026-07-26T13:35:52.820743Z", "iopub.status.idle": "2026-07-26T13:35:52.864070Z", "shell.execute_reply": "2026-07-26T13:35:52.863636Z" } }, "outputs": [ { "data": { "text/plain": [ "(Array([1.7421458, 1.7421458, 1.7421458], dtype=float32),)" ] }, "execution_count": 6, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# now we can compute the \"pullback\" of w through f at x;\n", "# returns a tuple, even for a single input,\n", "# since again f is assumed to take a \"product space of inputs\" *always*\n", "df_back(w)" ] }, { "cell_type": "markdown", "id": "22e94350", "metadata": {}, "source": [ "Note that we didn't need to specify $x$ when evaluating `df_back`.\n", "That is because `vjp` already evaluated $f$ in the \"forward\" direction and \"secretly\" stored all intermediate values it computed inside of `df_back` for later use.\n", "\n", "In terms of the python data types, for jax/numpy everything is just real numbers (or complex, see below).\n", "No distinction is made between $\\mathbb R^n$, $T\\mathbb R^n$ and $T^*\\mathbb R^n$.\n", "In particular, note that both cotangents and tangents are represented as \"row\" vectors of shape `(n,)`.\n", "To avoid confusion, it's also worth noting that it works exactly the same for any array shape, so for \"array\" shaped inputs, inputs, tangents and cotangents are all `(n, m)` shaped numpy arrays." ] }, { "cell_type": "markdown", "id": "f4c73061", "metadata": {}, "source": [ "### Grad\n", "\n", "`jax.grad` is just a convenience wrapper around the \"pullback\" `jax.vjp`.\n", "\n", "The largest number of ML users just want to take the gradient of the loss function -- a scalar function $l: \\Theta \\to L$ on the neural network parameters $\\theta \\in \\Theta$. Of course $\\Theta = \\mathbb{R}^{\\cdots}$ and $L=\\mathbb R$.\n", "Since parameters are just real numbers, carrying the standard (flat) metric, we can immediately identify the gradient $\\nabla l$ with the cotangent $1 \\circ dl \\in T^*\\Theta$.\n", "To be extra verbose, the last expression includes the cotangent $1 \\in T^*L$.\n", "This is exactly what the \"convenience\" function `jax.grad` implements." ] }, { "cell_type": "code", "execution_count": 7, "id": "d0b11ca0", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:52.865254Z", "iopub.status.busy": "2026-07-26T13:35:52.865184Z", "iopub.status.idle": "2026-07-26T13:35:52.925589Z", "shell.execute_reply": "2026-07-26T13:35:52.925169Z" } }, "outputs": [ { "data": { "text/plain": [ "Array([[2., 4.],\n", " [6., 8.]], dtype=float32)" ] }, "execution_count": 7, "metadata": {}, "output_type": "execute_result" } ], "source": [ "l = lambda theta: jnp.sum(theta**2)\n", "\n", "theta = jnp.array([[1.0, 2.0], [3.0, 4.0]])\n", "grad_fn = jax.grad(l)\n", "\n", "grad_fn(theta) # 2 * theta" ] }, { "cell_type": "code", "execution_count": 8, "id": "56486f41", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:52.926831Z", "iopub.status.busy": "2026-07-26T13:35:52.926752Z", "iopub.status.idle": "2026-07-26T13:35:52.950190Z", "shell.execute_reply": "2026-07-26T13:35:52.949849Z" }, "lines_to_next_cell": 1 }, "outputs": [ { "data": { "text/plain": [ "(Array([[2., 4.],\n", " [6., 8.]], dtype=float32),)" ] }, "execution_count": 8, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# jax.grad does a bit more book keeping, but the above is equivalent to:\n", "def our_grad(f):\n", " # just like grad, return a *function* that computes the gradient\n", "\n", " def grad_fn(x):\n", " f_x, df_back = jax.vjp(f, x)\n", " # could also return f(x) -- exactly jax.value_and_grad\n", " return df_back(1.0)\n", "\n", " return grad_fn\n", "\n", "# part of that book keeping: grad unwraps the one-element tuple, we don't\n", "our_grad(l)(theta)" ] }, { "cell_type": "markdown", "id": "87751c52", "metadata": {}, "source": [ "### Jacobians\n", "\n", "`jax.jacfwd` and `jax.jacrev` add no new capability either: they call `jvp` or `vjp` repeatedly, vectorized (`vmap`) over a *canonical basis*, to manifest the whole Jacobian as an array.\n", "\n", "- `jacfwd` pushes each basis **tangent** $e_j \\in TM$ forward, giving one *column* $df|_x(e_j)$ at a time.\n", "- `jacrev` pulls each basis **cotangent** $dy^i \\in T^*N$ back, giving one *row* $dy^i \\circ df|_x$ at a time.\n", "\n", "This is where the cost asymmetry comes from. For $f: \\mathbb R^n \\to \\mathbb R^m$, forward mode needs $n$ passes (one per *input* dimension) and reverse mode needs $m$ passes (one per *output* dimension). Neither is universally better; reverse mode wins when there are many inputs and few outputs, which is another reason why `grad` (where $m = 1$) is naturally reverse mode." ] }, { "cell_type": "code", "execution_count": 9, "id": "cd4669e2", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:52.951277Z", "iopub.status.busy": "2026-07-26T13:35:52.951208Z", "iopub.status.idle": "2026-07-26T13:35:53.186198Z", "shell.execute_reply": "2026-07-26T13:35:53.185829Z" } }, "outputs": [ { "data": { "text/plain": [ "(Array(True, dtype=bool), Array(True, dtype=bool))" ] }, "execution_count": 9, "metadata": {}, "output_type": "execute_result" } ], "source": [ "def our_jacfwd(f, x):\n", " basis = jnp.eye(x.shape[0]) # canonical basis of tangents\n", " push = lambda v: jax.jvp(f, (x,), (v,))[1]\n", " return jax.vmap(push)(basis).T # columns = pushed basis tangents\n", "\n", "def our_jacrev(f, x):\n", " y, f_back = jax.vjp(f, x)\n", " basis = jnp.eye(y.shape[0]) # canonical basis of cotangents\n", " pull = lambda w: f_back(w)[0]\n", " return jax.vmap(pull)(basis) # rows = pulled basis cotangents\n", "\n", "# f: R^3 -> R^2 from above, so this costs 3 forward passes or 2 reverse passes\n", "jnp.allclose(our_jacfwd(f, x), jax.jacfwd(f)(x)), jnp.allclose(our_jacrev(f, x), jax.jacrev(f)(x))" ] }, { "cell_type": "markdown", "id": "d36bced9", "metadata": {}, "source": [ "### Complex numbers\n", "\n", "In terms of vjp and jvp, jax treats $z \\in \\mathbb C$ identical to the vector `[x, y]` $\\in \\mathbb R^2$ such that $z = x + i y$.\n", "The function we differentiate through need not be holomorphic; it does what we would expect, thinking of $\\mathbb C$ as a Riemannian manifold.\n", "\n", "Concretely, decompose $z = x + iy$.\n", "Then $df = \\frac{\\partial f}{\\partial x} \\, dx + \\frac{\\partial f}{\\partial y} \\, dy$ or equivalently $df = \\frac{\\partial f}{\\partial z} \\, dz + \\frac{\\partial f}{\\partial \\bar z} \\, d\\bar{z}$.\n", "The mapping to how jax represents tangents and cotangents is straightforward, with the important note that the cotangent $v^\\dagger$ is represented as $v$ (not as \"$v^*$\").\n", "\n", "- $dx \\cong$ `1.0`\n", "- $dy \\cong$ `1.0j`\n", "- $dz \\cong$ `1.0 + 1.0j`\n", "- $d\\bar{z} \\cong$ `1.0 - 1.0j`\n", "- $\\partial_x \\cong$ `1.0`\n", "- $\\partial_y \\cong$ `1.0j`\n", "\n", "Note that it might be tempting to expect to be able to read off $\\partial_z f$ with a single `jvp` pass, by somehow specifying the tangent $\\partial_z = (\\partial_x - i \\partial_y)/2$.\n", "However, the latter is a member of the *complexified* tangent space of which JAX knows nothing. Just looking at dimensions we see it can't possibly work since the tangent space of $\\mathbb C$, for jax, is two dimensional but $\\partial_x, i\\partial_x, \\partial_y, i\\partial_y$ are distinct objects in the complexified tangent space. In particular, if $\\partial_y \\cong 1.0j$ then $i\\partial_y \\cong -1.0$ but the latter is already the same as $-\\partial_x$.\n", "\n", "We thus have to implement the $\\mathbb C$-linearity ourselves: $df(\\partial_z) = df(\\partial_x)/2 - i \\, df(\\partial_y)/2 \\cong$ `jvp(f, z, 1.0)/2 - 1j * jvp(f, z, 1.0j)`.\n" ] }, { "cell_type": "code", "execution_count": 10, "id": "2d9b1761", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:53.187434Z", "iopub.status.busy": "2026-07-26T13:35:53.187361Z", "iopub.status.idle": "2026-07-26T13:35:53.225128Z", "shell.execute_reply": "2026-07-26T13:35:53.224651Z" } }, "outputs": [ { "data": { "text/plain": [ "Array(1.+0.j, dtype=complex64, weak_type=True)" ] }, "execution_count": 10, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# consider complex conjugation\n", "conj = lambda z: z.real - 1j * z.imag # = z.conj()\n", "\n", "# note that z has to be complex if v is complex & vice versa\n", "# removing one 0.0j yields a type error\n", "z = 1.0 + 0.0j\n", "v = 1.0 + 0.0j\n", "\n", "out, out_tang = jax.jvp(conj, (z,), (v,))\n", "out_tang # directional derivative along the real axis gives 1" ] }, { "cell_type": "code", "execution_count": 11, "id": "e3315eb1", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:53.226149Z", "iopub.status.busy": "2026-07-26T13:35:53.226087Z", "iopub.status.idle": "2026-07-26T13:35:53.228756Z", "shell.execute_reply": "2026-07-26T13:35:53.228427Z" } }, "outputs": [ { "data": { "text/plain": [ "Array(0.-1.j, dtype=complex64, weak_type=True)" ] }, "execution_count": 11, "metadata": {}, "output_type": "execute_result" } ], "source": [ "z = 1.0 + 0.0j\n", "v = 0.0 + 1.0j\n", "\n", "out, out_tang = jax.jvp(conj, (z,), (v,))\n", "out_tang # directional derivative along the imaginary axis gives -i" ] }, { "cell_type": "markdown", "id": "683a48fb", "metadata": {}, "source": [ "We can now implement the derivative for complexified tangent vectors (either taking complex coefficients of $c_x \\partial_x + c_y \\partial_y$ or equivalently $c_z \\partial_z + c_{\\bar z} \\partial_{\\bar z}$ as inputs)." ] }, { "cell_type": "code", "execution_count": 12, "id": "3176a090", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:53.229755Z", "iopub.status.busy": "2026-07-26T13:35:53.229691Z", "iopub.status.idle": "2026-07-26T13:35:53.231382Z", "shell.execute_reply": "2026-07-26T13:35:53.230997Z" } }, "outputs": [], "source": [ "def complexified_deriv(f, z, coeff_x, coeff_y):\n", " \"\"\"Take derivative given complexified tangent vector.\"\"\"\n", " # could also define equivalent function with coeff_z, coeff_zbar as inputs\n", "\n", " # evaluate real JVPs for \\partial_x and \\partial_y\n", " df_dx = jax.jvp(f, (z,), (1.0 + 0.0j,))[1]\n", " df_dy = jax.jvp(f, (z,), (0.0 + 1.0j,))[1]\n", "\n", "\n", " # combine by C-linearity\n", " return coeff_x * df_dx + coeff_y * df_dy" ] }, { "cell_type": "code", "execution_count": 13, "id": "4d9f2d83", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:53.232403Z", "iopub.status.busy": "2026-07-26T13:35:53.232334Z", "iopub.status.idle": "2026-07-26T13:35:53.234316Z", "shell.execute_reply": "2026-07-26T13:35:53.233965Z" } }, "outputs": [ { "data": { "text/plain": [ "(1+0j)" ] }, "execution_count": 13, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# partial_z (z) == 1, as it should\n", "complexified_deriv(lambda z: z, 1.0 + 0j, 0.5, -0.5j)" ] }, { "cell_type": "markdown", "id": "217e2178", "metadata": {}, "source": [ "In particular we can now check if $f$ is holomorphic via the Cauchy-Riemann equation $\\partial_{\\bar z} f = 0$." ] }, { "cell_type": "code", "execution_count": 14, "id": "852ccde1", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:53.235204Z", "iopub.status.busy": "2026-07-26T13:35:53.235134Z", "iopub.status.idle": "2026-07-26T13:35:53.283730Z", "shell.execute_reply": "2026-07-26T13:35:53.283349Z" } }, "outputs": [ { "data": { "text/plain": [ "Array(0.+0.j, dtype=complex64, weak_type=True)" ] }, "execution_count": 14, "metadata": {}, "output_type": "execute_result" } ], "source": [ "partial_z = lambda f, z: complexified_deriv(f, z, 0.5, -0.5j)\n", "partial_zbar = lambda f, z: complexified_deriv(f, z, 0.5, 0.5j)\n", "\n", "partial_zbar(jnp.sin, 1.0j) # is holomorphic, so get 0" ] }, { "cell_type": "code", "execution_count": 15, "id": "4919994e", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:53.284969Z", "iopub.status.busy": "2026-07-26T13:35:53.284897Z", "iopub.status.idle": "2026-07-26T13:35:53.296915Z", "shell.execute_reply": "2026-07-26T13:35:53.296473Z" }, "lines_to_next_cell": 2 }, "outputs": [ { "data": { "text/plain": [ "Array(1.+0.j, dtype=complex64, weak_type=True)" ] }, "execution_count": 15, "metadata": {}, "output_type": "execute_result" } ], "source": [ "partial_zbar(jnp.conj, 1.0j) # not holomorphic, don't get 0 (in fact is antiholomorphic, so get 1)" ] }, { "cell_type": "markdown", "id": "e7ecfdaa", "metadata": {}, "source": [ "If we *know* a function is holomorphic, then Cauchy-Riemann $\\partial_{\\bar z} f = (\\partial_x + i \\partial_y) f / 2 = 0$ gives us $\\partial_y f = i \\, \\partial_x f$.\n", "Thus, we could save on one of the `jvp` calls: `df_dy = 1j * df_dx`.\n", "We could thus in principle define two holomorphic derivatives, one \"safe\" one that works on any function, and one \"unsafe\" one that is cheaper but *assumes* f is holomorphic.\n", "\n", "The analog story holds for `vjp`. We can again define a general version that works for the complexified cotangent space, and an \"unsafe\" one that assumes `f` is holomorphic." ] }, { "cell_type": "code", "execution_count": 16, "id": "c81ffba0", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:53.297968Z", "iopub.status.busy": "2026-07-26T13:35:53.297901Z", "iopub.status.idle": "2026-07-26T13:35:53.351334Z", "shell.execute_reply": "2026-07-26T13:35:53.350941Z" }, "lines_to_next_cell": 1 }, "outputs": [ { "data": { "text/plain": [ "(Array(1.6800001+2.99j, dtype=complex64, weak_type=True),\n", " Array(1.6800001+2.99j, dtype=complex64, weak_type=True))" ] }, "execution_count": 16, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# check the shortcut on a holomorphic function, before relying on it\n", "holo = lambda z: z**3 + 2 * z\n", "df_dx = jax.jvp(holo, (0.7 - 0.4j,), (1.0 + 0.0j,))[1]\n", "df_dy = jax.jvp(holo, (0.7 - 0.4j,), (0.0 + 1.0j,))[1]\n", "\n", "df_dy, 1j * df_dx # equal, so one jvp would have sufficed" ] }, { "cell_type": "code", "execution_count": 17, "id": "a177bdb0", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:53.352398Z", "iopub.status.busy": "2026-07-26T13:35:53.352329Z", "iopub.status.idle": "2026-07-26T13:35:53.354086Z", "shell.execute_reply": "2026-07-26T13:35:53.353796Z" } }, "outputs": [], "source": [ "def complexified_cotangent_deriv(f, z, coeff_dx, coeff_dy):\n", " \"\"\" Evaluates the pull-back of a complexified cotangent.\"\"\"\n", "\n", " _, vjp_fun = jax.vjp(f, z)\n", "\n", " # evaluate VJPs along standard real dual basis\n", " pullback_dx = vjp_fun(1.0 + 0.0j)[0] # cotangent w = 1, real part pull-back\n", " pullback_dy = vjp_fun(0.0 + 1.0j)[0] # cotangent w = i, imag part pull-back\n", "\n", "\n", " # combine by C-linearity\n", " return coeff_dx * pullback_dx + coeff_dy * pullback_dy" ] }, { "cell_type": "markdown", "id": "5c72b18d", "metadata": {}, "source": [ "For scalar functions, we may of course again want to apply `jax.grad`.\n", "If the function outputs real numbers and the input is complex, everything works as expected.\n", "If the output is complex, however, it's ambiguous which cotangent (in the two-dimensional real tangent space) to start with.\n", "For a holomorphic function $dy \\circ df = i \\, dx \\circ df$ (where now $x, y$ decompose the *output space* $\\mathbb C$), so we can just compute $dx \\circ df$ and have all information.\n", "This is why `jax.grad` has an optional argument `holomorphic` and if set `True` it applies `vjp` starting with `1.0`.\n", "For a multi-variable holomorphic function $f([z_1, z_2, z_3]) \\in \\mathbb C$, `jax.grad(f, holomorphic=True)` thus returns exactly the expected $z \\mapsto [df/dz_1, df/dz_2, df/dz_3]$, assuming $f$ is holomorphic (else writing the complex derivative $df/dz_1$ is nonsense)." ] }, { "cell_type": "markdown", "id": "5c67c615", "metadata": {}, "source": [ "#### Splitting trick\n", "\n", "There is one trick worth mentioning. If we know we have some non-holomorphic function $f$ of which we'll want to take holomorphic and antiholomorphic derivatives, we can write it as a function of two arguments, $f(z, \\bar{z})$.\n", "If we treat these two inputs as formally independent, and we promise that in the implementation of $f$ we only apply holomorphic functions in each of them, we can take the true (anti-) holomorphic derivative in a single backward/forward pass." ] }, { "cell_type": "markdown", "id": "9f34350a", "metadata": {}, "source": [ "# Exterior derivative operator\n", "\n", "Given the number of quantities involved, there are many possible \"book-keepings\" that don't change the math but change what we specify when and what we get out when.\n", "Some choices are suggested by efficiency.\n", "For example, we could compute the function value and the pushforward of a tangent vector in two separate passes, but if we *know* we want both that would be inefficient:\n", "the tangent propagation has to evaluate the forward chain of computations anyway.\n", "Nonetheless, different applications call for different computational structure, and the exterior derivative is a good place to see that play out: below we build it in forward mode, then try the obvious reverse-mode alternative and watch it fail to compose.\n", "\n", "The convention here: a $k$-form is a function of a point and $k$ tangent vectors, `omega(x, v_1, ..., v_k)`, returning its contraction. A $0$-form is then just a function of a point.\n", "For constant vector fields the exterior derivative is\n", "$$\n", "d\\omega(v_0, \\dots, v_k) = \\sum_i (-1)^i \\, \\partial_{v_i} \\, \\omega(v_0, \\dots, \\widehat{v_i}, \\dots, v_k) \\,,\n", "$$\n", "where $\\widehat{v_i}$ means that argument is omitted -- so each term is one `jvp` of $\\omega$ in the direction $v_i$, with the remaining tangents held fixed.\n", "\n", "Note that the below is in many ways not the most efficient implementation (e.g. replacing the loop-reduce with a single `vmap`ed call would be more efficient, keeping track if something was a total derivative would let us avoid computing $0$, ...)" ] }, { "cell_type": "code", "execution_count": 18, "id": "fc0cf8c7", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:53.355090Z", "iopub.status.busy": "2026-07-26T13:35:53.355023Z", "iopub.status.idle": "2026-07-26T13:35:53.356895Z", "shell.execute_reply": "2026-07-26T13:35:53.356594Z" } }, "outputs": [], "source": [ "from functools import reduce\n", "\n", "# helper: choose one out of n\n", "def _pick_one(elements):\n", " parity = 1\n", " for i in range(len(elements)):\n", " yield parity, elements[i], (*elements[:i], *elements[i + 1:])\n", " parity *= -1\n", "\n", "# take exterior derivative\n", "def extd(fn):\n", "\n", " def d_fn(x, *tangents):\n", " return reduce(jnp.add, [\n", " jax.jvp(\n", " # want to take derivative only in x, so must hide ts dependence\n", " lambda _x: p * fn(_x, *ts),\n", " (x,), (t,)\n", " )[1]\n", " for p, t, ts in _pick_one(tangents)\n", " ])\n", "\n", " return d_fn" ] }, { "cell_type": "code", "execution_count": 19, "id": "908ec02d", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:53.357746Z", "iopub.status.busy": "2026-07-26T13:35:53.357696Z", "iopub.status.idle": "2026-07-26T13:35:53.400825Z", "shell.execute_reply": "2026-07-26T13:35:53.400483Z" } }, "outputs": [ { "data": { "text/plain": [ "Array(9., dtype=float32, weak_type=True)" ] }, "execution_count": 19, "metadata": {}, "output_type": "execute_result" } ], "source": [ "cube = lambda x: x**3\n", "d_cube = extd(cube) # d_cube(x, t) = 3 * x^2 * t\n", "\n", "d_cube(3.0, 1/3)" ] }, { "cell_type": "markdown", "id": "337baf6c", "metadata": {}, "source": [ "Next, $d^2 = 0$. Note we have to leave $\\mathbb R^1$ to test this properly: a $2$-form on a one-dimensional space vanishes for trivial reasons ($\\Lambda^2 \\mathbb R^1 = 0$), so it wouldn't really test the cancellation. On $\\mathbb R^2$ the two terms genuinely have to cancel." ] }, { "cell_type": "code", "execution_count": 20, "id": "23e06b6b", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:53.401971Z", "iopub.status.busy": "2026-07-26T13:35:53.401909Z", "iopub.status.idle": "2026-07-26T13:35:53.514066Z", "shell.execute_reply": "2026-07-26T13:35:53.513542Z" } }, "outputs": [ { "data": { "text/plain": [ "Array(0., dtype=float32)" ] }, "execution_count": 20, "metadata": {}, "output_type": "execute_result" } ], "source": [ "h = lambda x: jnp.sin(x[0]) * x[1]**2 # a 0-form on R^2\n", "ddh = extd(extd(h))\n", "\n", "e1, e2 = jnp.array([1., 0.]), jnp.array([0., 1.])\n", "ddh(jnp.array([0.3, -0.7]), e1, e2)" ] }, { "cell_type": "markdown", "id": "9cf7a640", "metadata": {}, "source": [ "As a slightly less trivial example, consider $g(x) = x^1\\,dx^2 - x^2\\,dx^1$. We might be tempted to immediately identify this with the (tangent) vector\n", "$$\n", "\\begin{pmatrix}\n", "-x^2 \\\\\n", "x^1\n", "\\end{pmatrix}\n", "= \\nabla \\times r^2/2 = \\nabla \\times ((x^1)^2 + (x^2)^2)/2\n", "$$\n", "and thus think we should implement $g$ as a map $\\mathbb R^2 \\to \\mathbb R^2$.\n", "But that doesn't quite fit the above, where we picked `jvp` and thus need to take tangent vectors as input.\n", "Thus we should implement $g$ as a map $\\mathbb R^2 \\to (T \\mathbb R^2 \\to \\mathbb R)$ which implements the contraction $g(x) \\cdot v$." ] }, { "cell_type": "code", "execution_count": 21, "id": "2028c70d", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:53.515177Z", "iopub.status.busy": "2026-07-26T13:35:53.515091Z", "iopub.status.idle": "2026-07-26T13:35:53.539142Z", "shell.execute_reply": "2026-07-26T13:35:53.538465Z" } }, "outputs": [ { "data": { "text/plain": [ "Array(-1., dtype=float32)" ] }, "execution_count": 21, "metadata": {}, "output_type": "execute_result" } ], "source": [ "g = lambda x, v: v[1] * x[0] - v[0] * x[1]\n", "\n", "g(jnp.array([0., 1.]), jnp.array([1., 0.])) # -1" ] }, { "cell_type": "markdown", "id": "94368192", "metadata": {}, "source": [ "Now $g$ is not closed, $dg = 2 \\, dx^1 \\wedge dx^2$ so we can test `extd` on 1-forms." ] }, { "cell_type": "code", "execution_count": 22, "id": "95986d3b", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:53.540706Z", "iopub.status.busy": "2026-07-26T13:35:53.540593Z", "iopub.status.idle": "2026-07-26T13:35:53.567093Z", "shell.execute_reply": "2026-07-26T13:35:53.566597Z" } }, "outputs": [ { "data": { "text/plain": [ "Array(2., dtype=float32)" ] }, "execution_count": 22, "metadata": {}, "output_type": "execute_result" } ], "source": [ "dg = extd(g)\n", "\n", "# g is not closed, so this is not always zero as ddh above was\n", "dg(jnp.array([1., 1.]), e1, e2)" ] }, { "cell_type": "code", "execution_count": 23, "id": "c6a1635f", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:53.568274Z", "iopub.status.busy": "2026-07-26T13:35:53.568197Z", "iopub.status.idle": "2026-07-26T13:35:53.571853Z", "shell.execute_reply": "2026-07-26T13:35:53.571413Z" } }, "outputs": [ { "data": { "text/plain": [ "Array(-2., dtype=float32)" ] }, "execution_count": 23, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# check orientation flip gives -1\n", "dg(jnp.array([1., 1.]), e2, e1)" ] }, { "cell_type": "code", "execution_count": 24, "id": "bf07e66c", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:53.573329Z", "iopub.status.busy": "2026-07-26T13:35:53.573240Z", "iopub.status.idle": "2026-07-26T13:35:53.576830Z", "shell.execute_reply": "2026-07-26T13:35:53.576453Z" } }, "outputs": [ { "data": { "text/plain": [ "Array(0., dtype=float32)" ] }, "execution_count": 24, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# collapse if tangents collinear\n", "dg(jnp.array([1., 1.]), jnp.array([1., 1.]), jnp.array([2., 2.]))" ] }, { "cell_type": "markdown", "id": "c9c1a230", "metadata": {}, "source": [ "### Why composition worked\n", "\n", "Worth thinking about briefly, because the `vjp` case below is less straightforward.\n", "We represented a $k$-form as a function of a point and $k$ tangent vectors, contracted down to a number:\n", "$$\n", "\\mathrm{Form}_k = M \\times (TM)^k \\to \\mathbb R \\,, \\qquad\n", "\\mathrm{extd}: \\mathrm{Form}_k \\to \\mathrm{Form}_{k+1} \\,.\n", "$$\n", "The output type is the input type with $k$ raised by one -- the representation is *closed* under the operation.\n", "That is what lets us iterate: `extd` returns the same kind of object it consumes, so it can be fed back in.\n", "Note this is a statement about our chosen book-keeping, not about the mathematics (see below)." ] }, { "cell_type": "markdown", "id": "991b81e2", "metadata": {}, "source": [ "### Reverse mode: the obvious thing to try\n", "\n", "The forward version makes us pay one `jvp` per tangent we want to contract against.\n", "For a $0$-form that looks wasteful: $df|_x$ has $n$ components and reverse mode is supposed to hand us all of them in a single pass.\n", "So the obvious move is to write `d` with `vjp` instead.\n", "\n", "Consider again $f: M \\to N$. If $N = \\mathbb R$ we can start from the \"canonical\" cotangent $1$; for general $N$ we have to say which cotangent on $N$ to project onto first." ] }, { "cell_type": "code", "execution_count": 25, "id": "0d0339e1", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:53.577822Z", "iopub.status.busy": "2026-07-26T13:35:53.577769Z", "iopub.status.idle": "2026-07-26T13:35:53.579381Z", "shell.execute_reply": "2026-07-26T13:35:53.578997Z" }, "lines_to_next_cell": 1 }, "outputs": [], "source": [ "def extd_rev(fn):\n", "\n", " def d_fn(x, cot=1.0):\n", " _, fn_back = jax.vjp(fn, x)\n", " (df_x,) = fn_back(cot)\n", " return df_x\n", "\n", " return d_fn" ] }, { "cell_type": "code", "execution_count": 26, "id": "51b5b461", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:53.580485Z", "iopub.status.busy": "2026-07-26T13:35:53.580404Z", "iopub.status.idle": "2026-07-26T13:35:53.654165Z", "shell.execute_reply": "2026-07-26T13:35:53.653714Z" } }, "outputs": [ { "data": { "text/plain": [ "Array([ 0.46811488, -0.4137283 ], dtype=float32)" ] }, "execution_count": 26, "metadata": {}, "output_type": "execute_result" } ], "source": [ "x0 = jnp.array([0.3, -0.7])\n", "\n", "# same object as extd(h), but as cotangent components instead of a contraction:\n", "# the forward version gives us one number per tangent we supply,\n", "# the reverse version gives the whole covector in one pass.\n", "extd_rev(h)(x0)" ] }, { "cell_type": "code", "execution_count": 27, "id": "aa4448ee", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:53.655908Z", "iopub.status.busy": "2026-07-26T13:35:53.655795Z", "iopub.status.idle": "2026-07-26T13:35:53.672864Z", "shell.execute_reply": "2026-07-26T13:35:53.672407Z" }, "lines_to_next_cell": 2 }, "outputs": [ { "data": { "text/plain": [ "(Array(0.46811488, dtype=float32), Array(-0.4137283, dtype=float32))" ] }, "execution_count": 27, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# can extract both with forward mode by projecting onto bases\n", "(extd(h)(x0, e1), extd(h)(x0, e2))" ] }, { "cell_type": "markdown", "id": "bab66007", "metadata": {}, "source": [ "...but it only works once. If we try to iterate:" ] }, { "cell_type": "code", "execution_count": 28, "id": "4bb00daf", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:53.674246Z", "iopub.status.busy": "2026-07-26T13:35:53.674160Z", "iopub.status.idle": "2026-07-26T13:35:53.710278Z", "shell.execute_reply": "2026-07-26T13:35:53.709857Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "ValueError: unexpected JAX type (e.g. shape/dtype) for argument to VJP function: got float32[], but expected float32[2] because the corresponding output of the differentiated function had JAX type float32[2]\n" ] } ], "source": [ "try:\n", " extd_rev(extd_rev(h))(x0)\n", "except ValueError as err:\n", " print('ValueError:', err)" ] }, { "cell_type": "markdown", "id": "c741645c", "metadata": {}, "source": [ "Two things go wrong.\n", "\n", "**The direction flips.** \n", "`jvp` always runs input $\\to$ output, so pushing forward again just continues in the same direction. But `vjp` returns a map *into* $T^*M$, cotangents on the **input** space. So `omega = extd_rev(h)` is a function $M \\to T^*M$, and differentiating *that* in reverse mode needs a seed which pairs with a covector, i.e. a **tangent**, not a cotangent. One more level and it flips back. The seed type alternates $T^*, T, T^*, $ etc.\n", "\n", "Also, note that the `ValueError` we did get is only the second point showing up by luck.\n", "In principle `vjp` could run another time, but only if we specify a projection cotangent, now of the shape of the space.\n", "That is because it makes no distinction between $\\mathbb R^n$, $T\\mathbb R^n$ and $T^*\\mathbb R^n$, and so interprets the map generated by the first `extd_rev` as simply a map from $\\mathbb R^n$ to $\\mathbb R^n$. The bookkeeping is our responsibility." ] }, { "cell_type": "markdown", "id": "b1e23081", "metadata": {}, "source": [ "We can of course still build the $2$-form in reverse mode, by doing the $n$ passes explicitly and antisymmetrizing. Which is precisely `jacrev`: `vjp` vectorized over a basis." ] }, { "cell_type": "code", "execution_count": 29, "id": "3373e572", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:53.711562Z", "iopub.status.busy": "2026-07-26T13:35:53.711468Z", "iopub.status.idle": "2026-07-26T13:35:53.860464Z", "shell.execute_reply": "2026-07-26T13:35:53.860028Z" }, "lines_to_next_cell": 2 }, "outputs": [ { "data": { "text/plain": [ "Array([[ 0. , 2.1 , 0.25533652],\n", " [-2.1 , 0. , 0. ],\n", " [-0.25533652, 0. , 0. ]], dtype=float32)" ] }, "execution_count": 29, "metadata": {}, "output_type": "execute_result" } ], "source": [ "n = 3\n", "basis = jnp.eye(n)\n", "omega = lambda x: jnp.array([-x[1] * x[2], x[0], jnp.sin(x[0])])\n", "x1 = jnp.array([0.3, -0.7, 1.1])\n", "\n", "# one reverse pass per output component, with the flipped-type seed\n", "slices = jnp.stack([extd_rev(omega)(x1, basis[j]) for j in range(n)]) # slices[j, i] = d_i omega_j\n", "slices.T - slices # (d omega)_{ij} = d_i omega_j - d_j omega_i" ] }, { "cell_type": "markdown", "id": "532efb7e", "metadata": {}, "source": [ "### Closing the type again: components instead of contractions\n", "\n", "The fix is not to patch `extd_rev` but to change the representation. Represent a $k$-form by its **component array** rather than by its contraction,\n", "$$\n", "\\mathrm{Form}_k = M \\to \\Lambda^k T^*M \\,,\n", "$$\n", "i.e. a function from a point to an antisymmetric array of shape `(n,) * k`. Now $d$ is one Jacobian plus an antisymmetrization,\n", "$$\n", "(d\\omega)_{i_0 \\dots i_k} = (k+1) \\, \\mathrm{Antisym} \\left( \\partial_{i_0} \\omega_{i_1 \\dots i_k} \\right) \\,,\n", "$$\n", "and the output is again a component array, so the type closes and we can iterate.\n", "Note this version is agnostic about the mode: the Jacobian can come from `jacfwd` or `jacrev` and the answer is identical.\n", "\n", "This representation does demand something the contraction version did not: the point has to be a shape-`(n,)` array, `[0.1]` rather than `0.1`, even when $n = 1$.\n", "The reason is that the array axes *are* the form indices, so each derivative needs an axis of the point to hang its new index on, and a bare scalar has none.\n", "It is worth being strict about this rather than quietly calling `atleast_1d` on the Jacobian: that would make the one-dimensional case run, but on the second application there is still no input axis to distinguish from the form index, and the two collapse.\n", "The result is a plain second derivative instead of an antisymmetrization -- for $x^3$ at $x = 0.1$ it returns $0.6$ where $d^2 = 0$ demands $0$.\n", "A wrong answer is worse than a failure, so we assert the shape instead.\n", "The contraction-based `extd`, by contrast, never indexes anything and so is indifferent to the encoding: `d_cube(3.0, 1/3)` above passed a bare float quite happily." ] }, { "cell_type": "code", "execution_count": 30, "id": "1d25ae1c", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:53.861708Z", "iopub.status.busy": "2026-07-26T13:35:53.861628Z", "iopub.status.idle": "2026-07-26T13:35:53.864174Z", "shell.execute_reply": "2026-07-26T13:35:53.863771Z" }, "lines_to_next_cell": 1 }, "outputs": [], "source": [ "import math\n", "from itertools import permutations\n", "\n", "def _parity(perm):\n", " \"\"\"Sign of a permutation, by counting inversions.\"\"\"\n", " p, sgn = list(perm), 1\n", " for i in range(len(p)):\n", " for j in range(i + 1, len(p)):\n", " if p[i] > p[j]:\n", " sgn = -sgn\n", " return sgn\n", "\n", "def antisym(arr):\n", " \"\"\"Normalized antisymmetrization over all axes.\"\"\"\n", " terms = [_parity(p) * jnp.transpose(arr, p) for p in permutations(range(arr.ndim))]\n", " return reduce(jnp.add, terms) / math.factorial(arr.ndim)\n", "\n", "def extd_comp(omega, jac=jax.jacfwd):\n", " \"\"\"Exterior derivative of a k-form given by components omega: M -> Lambda^k T*M.\n", "\n", " The point must be a shape-(n,) array -- [0.1], not 0.1 -- even for n = 1.\n", " \"\"\"\n", "\n", " def d_omega(x):\n", " assert jnp.ndim(x) == 1, 'point must be a shape-(n,) array, not a bare scalar'\n", " # jac(omega)(x) has shape (n,) * k + (n,); move the input index to the front\n", " arr = jnp.moveaxis(jac(omega)(x), -1, 0)\n", " return arr.ndim * antisym(arr) # arr.ndim == k + 1\n", "\n", " return d_omega" ] }, { "cell_type": "code", "execution_count": 31, "id": "205dd60d", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:53.865410Z", "iopub.status.busy": "2026-07-26T13:35:53.865326Z", "iopub.status.idle": "2026-07-26T13:35:53.980663Z", "shell.execute_reply": "2026-07-26T13:35:53.979967Z" } }, "outputs": [ { "data": { "text/plain": [ "Array([[ 0. , 2.1 , 0.25533652],\n", " [-2.1 , 0. , 0. ],\n", " [-0.25533652, 0. , 0. ]], dtype=float32)" ] }, "execution_count": 31, "metadata": {}, "output_type": "execute_result" } ], "source": [ "extd_comp(omega)(x1)" ] }, { "cell_type": "code", "execution_count": 32, "id": "f2f09a33", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:53.982018Z", "iopub.status.busy": "2026-07-26T13:35:53.981913Z", "iopub.status.idle": "2026-07-26T13:35:54.027907Z", "shell.execute_reply": "2026-07-26T13:35:54.027394Z" } }, "outputs": [ { "data": { "text/plain": [ "(Array(2.1, dtype=float32), Array(2.1, dtype=float32))" ] }, "execution_count": 32, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# cross-check the two representations against each other: contract omega by hand,\n", "# hand it to the forward extd, and compare with one entry of the component version\n", "omega_contracted = lambda x, v: jnp.dot(omega(x), v)\n", "\n", "extd(omega_contracted)(x1, basis[0], basis[1]), extd_comp(omega)(x1)[0, 1]" ] }, { "cell_type": "code", "execution_count": 33, "id": "1479d821", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:54.029169Z", "iopub.status.busy": "2026-07-26T13:35:54.029095Z", "iopub.status.idle": "2026-07-26T13:35:54.602505Z", "shell.execute_reply": "2026-07-26T13:35:54.601984Z" } }, "outputs": [ { "data": { "text/plain": [ "(Array(0., dtype=float32), Array(0., dtype=float32))" ] }, "execution_count": 33, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# ...and now it composes, so d^2 = 0 is checkable on the manifest tensor\n", "extd_comp(extd_comp(omega))(x1).max(), extd_comp(extd_comp(h))(x0).max()" ] }, { "cell_type": "code", "execution_count": 34, "id": "a5f91036", "metadata": { "execution": { "iopub.execute_input": "2026-07-26T13:35:54.603706Z", "iopub.status.busy": "2026-07-26T13:35:54.603635Z", "iopub.status.idle": "2026-07-26T13:35:54.742202Z", "shell.execute_reply": "2026-07-26T13:35:54.741722Z" } }, "outputs": [ { "data": { "text/plain": [ "Array(True, dtype=bool)" ] }, "execution_count": 34, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# the mode genuinely does not matter for the result, only for the cost\n", "jnp.allclose(extd_comp(omega, jax.jacfwd)(x1), extd_comp(omega, jax.jacrev)(x1))" ] } ], "metadata": { "kernelspec": { "display_name": ".venv", "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.13.5" } }, "nbformat": 4, "nbformat_minor": 5 }