{ "cells": [ { "cell_type": "raw", "id": "49c84566", "metadata": {}, "source": [ "---\n", "Copyright 2026 The JAX Authors.\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", "ruff: noqa: F811\n", "\n", "---" ] }, { "cell_type": "markdown", "id": "b938e564", "metadata": {}, "source": [ "(hijax-custom-derivatives)=\n", "# Custom derivative rules with hijax primitives\n", "\n", "There have long been two ways to define differentiation rules in JAX:\n", "\n", "1. using `jax.custom_jvp` and `jax.custom_vjp` to define custom\n", " differentiation rules for Python functions that are already\n", " JAX-transformable, as described in [the custom derivatives\n", " notebook](https://docs.jax.dev/en/latest/notebooks/Custom_derivative_rules_for_Python_code.html);\n", " and\n", "2. defining new `core.Primitive` instances along with all their\n", " transformation rules, for example to call into functions from other\n", " systems like solvers, simulators, or general numerical computing systems.\n", "\n", "Hijax primitives unify these two. A hijax primitive is a real JAX primitive:\n", "it appears as a single unit in jaxprs, and you can attach a rule for each\n", "transformation. But like a `jax.custom_vjp`-decorated function, it also\n", "carries a Python implementation, written in terms of ordinary JAX operations,\n", "which is used for evaluation and lowering. So you don't need to write a\n", "lowering rule, and rules for the other transformations are only needed if you\n", "actually use those transformations.\n", "\n", "This document shows how to use hijax primitives to define custom derivative\n", "rules for JAX-transformable Python code. It's a hijax-flavored analogue of\n", "[the custom derivatives\n", "notebook](https://docs.jax.dev/en/latest/notebooks/Custom_derivative_rules_for_Python_code.html),\n", "and assumes familiarity with `jax.jvp` and `jax.grad`, and the mathematical\n", "meaning of JVPs and VJPs (see [The Autodiff\n", "Cookbook](https://docs.jax.dev/en/latest/notebooks/autodiff_cookbook.html)).\n", "\n", "Hijax is experimental: expect imports from `jax.experimental.hijax`, and\n", "expect the APIs to evolve. For stable APIs that cover most custom-derivative\n", "use cases, use `jax.custom_jvp` and `jax.custom_vjp`, described in\n", "{ref}`advanced-autodiff-custom-derivative-rules`.\n", "\n", "### TL;DR\n", "\n", "Subclass `HiPrim`, declare the input and output types, and define an\n", "`expand` method giving the implementation. Then attach the derivative rules\n", "you need: `vjp_fwd`/`vjp_bwd_retval` for reverse mode, `jvp` for forward\n", "mode:" ] }, { "cell_type": "code", "execution_count": null, "id": "e5ecfa03", "metadata": {}, "outputs": [], "source": [ "import jax\n", "import jax.numpy as jnp\n", "from jax.experimental.hijax import HiPrim\n", "\n", "class SinTimesY(HiPrim):\n", " def __init__(self, x_aval, y_aval):\n", " self.in_avals = (x_aval, y_aval) # input types\n", " self.out_aval = x_aval # output type\n", " self.params = {} # static parameters (none here)\n", " super().__init__()\n", "\n", " # Implementation, used for evaluation and lowering (e.g. under jit).\n", " def expand(self, x, y):\n", " return jnp.sin(x) * y\n", "\n", " # Reverse-mode: forward pass returns (primal_out, residuals).\n", " def vjp_fwd(self, nzs_in, x, y):\n", " return self(x, y), (jnp.cos(x), jnp.sin(x), y)\n", "\n", " # Reverse-mode: backward pass maps (residuals, output cotangent) to a tuple\n", " # of input cotangents.\n", " def vjp_bwd_retval(self, res, g):\n", " cos_x, sin_x, y = res\n", " return (cos_x * g * y, sin_x * g)\n", "\n", " # Forward-mode rule (optional, only needed under e.g. jax.jvp).\n", " def jvp(self, primals, tangents):\n", " (x, y), (x_dot, y_dot) = primals, tangents\n", " return self(x, y), jnp.cos(x) * x_dot * y + jnp.sin(x) * y_dot\n", "\n", "def f(x, y):\n", " return SinTimesY(jax.typeof(x), jax.typeof(y))(x, y)" ] }, { "cell_type": "code", "execution_count": null, "id": "ab0491e8", "metadata": {}, "outputs": [], "source": [ "from jax import jvp, grad\n", "\n", "print(f(2., 3.))\n", "y, y_dot = jvp(f, (2., 3.), (1., 0.))\n", "print(y)\n", "print(y_dot)\n", "print(grad(f)(2., 3.))" ] }, { "cell_type": "markdown", "id": "f1416ebf", "metadata": {}, "source": [ "Note that, unlike with `jax.custom_jvp` and `jax.custom_vjp` where you must\n", "choose one or the other, a single hijax primitive can carry both forward- and\n", "reverse-mode rules.\n", "\n", "## Example problems\n", "\n", "To get an idea of what problems hijax custom derivatives are meant to solve,\n", "let's go over a few examples. A more thorough introduction to the\n", "`HiPrim` API is in the next section.\n", "\n", "### Example: Numerical stability\n", "\n", "One application of custom derivative rules is to improve the numerical\n", "stability of differentiation.\n", "\n", "Say we want to write a function called `log1pexp`, which computes\n", "$x \\mapsto \\log ( 1 + e^x )$. We can write that using `jax.numpy`:" ] }, { "cell_type": "code", "execution_count": null, "id": "781e8f9c", "metadata": {}, "outputs": [], "source": [ "def log1pexp(x):\n", " return jnp.log(1. + jnp.exp(x))\n", "\n", "log1pexp(3.)" ] }, { "cell_type": "markdown", "id": "e566d500", "metadata": {}, "source": [ "Since it's written in terms of `jax.numpy`, it's JAX-transformable:" ] }, { "cell_type": "code", "execution_count": null, "id": "e8475e64", "metadata": {}, "outputs": [], "source": [ "from jax import jit, grad, vmap\n", "\n", "print(jit(log1pexp)(3.))\n", "print(jit(grad(log1pexp))(3.))\n", "print(vmap(jit(grad(log1pexp)))(jnp.arange(3.)))" ] }, { "cell_type": "markdown", "id": "38639dfd", "metadata": {}, "source": [ "But there's a numerical stability problem lurking here:" ] }, { "cell_type": "code", "execution_count": null, "id": "42fbd8e7", "metadata": {}, "outputs": [], "source": [ "print(grad(log1pexp)(100.))" ] }, { "cell_type": "markdown", "id": "1c005ea8", "metadata": {}, "source": [ "That doesn't seem right! After all, the derivative of $x \\mapsto \\log (1 +\n", "e^x)$ is $x \\mapsto \\frac{e^x}{1 + e^x}$, and so for large values of $x$ we'd\n", "expect the value to be about 1.\n", "\n", "We can get a bit more insight into what's going on by looking at the jaxpr\n", "for the gradient computation:" ] }, { "cell_type": "code", "execution_count": null, "id": "efa15290", "metadata": {}, "outputs": [], "source": [ "jit(grad(log1pexp)).trace(100.).jaxpr" ] }, { "cell_type": "markdown", "id": "53579928", "metadata": {}, "source": [ "Stepping through how the jaxpr would be evaluated, we can see that the last\n", "line would involve multiplying values that floating point math will round to\n", "0 and $\\infty$, respectively, which is never a good idea. That is, we're\n", "effectively evaluating `lambda x: (1 / (1 + jnp.exp(x))) * jnp.exp(x)` for\n", "large `x`, which effectively turns into `0. * jnp.inf`.\n", "\n", "Instead of generating such large and small values, hoping for a cancellation\n", "that floats can't always provide, we'd rather just express the derivative\n", "function as a more numerically stable program. In particular, we can write a\n", "program that more closely evaluates the equal mathematical expression $1 -\n", "\\frac{1}{1 + e^x}$, with no cancellation in sight.\n", "\n", "This problem is interesting because even though our definition of `log1pexp`\n", "could already be JAX-differentiated (and transformed with `jit`, `vmap`,\n", "...), we're not happy with the result of applying standard autodiff rules to\n", "the primitives comprising `log1pexp` and composing the result. Instead, we'd\n", "like to specify how the whole function `log1pexp` should be differentiated,\n", "as a unit, and thus arrange those exponentials better.\n", "\n", "Here's a solution using a hijax primitive:" ] }, { "cell_type": "code", "execution_count": null, "id": "74b21b10", "metadata": {}, "outputs": [], "source": [ "class Log1pExp(HiPrim):\n", " def __init__(self, x_aval):\n", " self.in_avals = (x_aval,)\n", " self.out_aval = x_aval\n", " self.params = {}\n", " super().__init__()\n", "\n", " def expand(self, x):\n", " return jnp.log(1. + jnp.exp(x))\n", "\n", " def vjp_fwd(self, nzs_in, x):\n", " return self(x), x\n", "\n", " def vjp_bwd_retval(self, x, g):\n", " return ((1 - 1/(1 + jnp.exp(x))) * g,)\n", "\n", " def jvp(self, primals, tangents):\n", " (x,), (x_dot,) = primals, tangents\n", " return self(x), (1 - 1/(1 + jnp.exp(x))) * x_dot\n", "\n", " def batch(self, axis_data, args, dims):\n", " (x,), (x_dim,) = args, dims\n", " return log1pexp(x), x_dim\n", "\n", "def log1pexp(x):\n", " return Log1pExp(jax.typeof(x))(x)" ] }, { "cell_type": "code", "execution_count": null, "id": "b3a5eb78", "metadata": {}, "outputs": [], "source": [ "print(grad(log1pexp)(100.))" ] }, { "cell_type": "code", "execution_count": null, "id": "896c9533", "metadata": {}, "outputs": [], "source": [ "print(jit(log1pexp)(3.))\n", "print(jit(grad(log1pexp))(3.))\n", "print(vmap(jit(grad(log1pexp)))(jnp.arange(3.)))" ] }, { "cell_type": "markdown", "id": "e79600f5", "metadata": {}, "source": [ "The `expand` method plays the role of the original Python function: it's\n", "what runs when we evaluate `log1pexp` eagerly, and it's what gets traced when\n", "we apply `jit`. The `vjp_fwd`/`vjp_bwd_retval` pair defines reverse-mode\n", "differentiation, and `jvp` defines forward-mode. (The `batch` method tells\n", "`vmap` how to batch the primitive: here it applies the primitive to the\n", "batched argument, whose batch dimension carries through. More on that below.)\n", "\n", "We can inspect the jaxpr of the gradient computation to confirm the stable\n", "formula is what runs:" ] }, { "cell_type": "code", "execution_count": null, "id": "b58ee782", "metadata": {}, "outputs": [], "source": [ "jit(grad(log1pexp)).trace(100.).jaxpr" ] }, { "cell_type": "markdown", "id": "306a21f2", "metadata": {}, "source": [ "### Example: Enforcing a differentiation convention\n", "\n", "A related application is to enforce a differentiation convention, perhaps at\n", "a boundary.\n", "\n", "Consider the function $f : \\mathbb{R}_+ \\to \\mathbb{R}_+$ with $f(x) =\n", "\\frac{x}{1 + \\sqrt{x}}$, where we take $\\mathbb{R}_+ = [0, \\infty)$. We might\n", "implement $f$ as a program like this:" ] }, { "cell_type": "code", "execution_count": null, "id": "cf381fc0", "metadata": {}, "outputs": [], "source": [ "def f(x):\n", " return x / (1 + jnp.sqrt(x))" ] }, { "cell_type": "markdown", "id": "12b5b18f", "metadata": {}, "source": [ "As a mathematical function on $\\mathbb{R}$ (the full real line), $f$ is not\n", "differentiable at zero (because the limit defining the derivative doesn't\n", "exist from the left). Correspondingly, autodiff produces a `nan` value:" ] }, { "cell_type": "code", "execution_count": null, "id": "9c8498d6", "metadata": {}, "outputs": [], "source": [ "print(grad(f)(0.))" ] }, { "cell_type": "markdown", "id": "01f7c526", "metadata": {}, "source": [ "But mathematically if we think of $f$ as a function on $\\mathbb{R}_+$ then it\n", "is differentiable at 0 [Rudin's Principles of Mathematical Analysis\n", "Definition 5.1, or Tao's Analysis I 3rd ed. Definition 10.1.1 and Example\n", "10.1.6]. Alternatively, we might say as a convention we want to consider the\n", "directional derivative from the right. So there is a sensible value for the\n", "Python function `grad(f)` to return at `0.0`, namely `1.0`. By default, JAX's\n", "machinery for differentiation assumes all functions are defined over\n", "$\\mathbb{R}$ and thus doesn't produce `1.0` here.\n", "\n", "We can use a custom derivative rule! In particular, we can define the rule in\n", "terms of the derivative function $x \\mapsto \\frac{\\sqrt{x} + 2}{2(\\sqrt{x} +\n", "1)^2}$ on $\\mathbb{R}_+$:" ] }, { "cell_type": "code", "execution_count": null, "id": "e80cd197", "metadata": {}, "outputs": [], "source": [ "class FOnRPlus(HiPrim):\n", " def __init__(self, x_aval):\n", " self.in_avals = (x_aval,)\n", " self.out_aval = x_aval\n", " self.params = {}\n", " super().__init__()\n", "\n", " def expand(self, x):\n", " return x / (1 + jnp.sqrt(x))\n", "\n", " def vjp_fwd(self, nzs_in, x):\n", " return self(x), x\n", "\n", " def vjp_bwd_retval(self, x, g):\n", " return (((jnp.sqrt(x) + 2) / (2 * (jnp.sqrt(x) + 1)**2)) * g,)\n", "\n", "def f(x):\n", " return FOnRPlus(jax.typeof(x))(x)" ] }, { "cell_type": "code", "execution_count": null, "id": "95111799", "metadata": {}, "outputs": [], "source": [ "print(grad(f)(0.))" ] }, { "cell_type": "markdown", "id": "dacb1e19", "metadata": {}, "source": [ "### Example: Gradient clipping\n", "\n", "While in some cases we want to express a mathematical differentiation\n", "computation, in other cases we may even want to take a step away from\n", "mathematics to adjust the computation autodiff performs. One canonical\n", "example is reverse-mode gradient clipping.\n", "\n", "For gradient clipping, we can use `jnp.clip` together with a\n", "reverse-mode-only rule. The bounds `lo` and `hi` are ordinary traced inputs:\n", "they aren't involved in differentiation, but they are runtime data (they\n", "could even be `jit` tracers), so the forward rule saves them as residuals\n", "and the backward rule returns `None` for their cotangents:" ] }, { "cell_type": "code", "execution_count": null, "id": "67390f76", "metadata": {}, "outputs": [], "source": [ "def batch_dims_at_front(axis_data, args, dims):\n", " return [jnp.moveaxis(a, d, 0) if d is not None\n", " else jnp.broadcast_to(a, (axis_data.size, *jnp.shape(a)))\n", " for a, d in zip(args, dims)]\n", "\n", "class ClipGradient(HiPrim):\n", " def __init__(self, lo_aval, hi_aval, x_aval):\n", " self.in_avals = (lo_aval, hi_aval, x_aval)\n", " self.out_aval = x_aval\n", " self.params = {}\n", " super().__init__()\n", "\n", " def expand(self, lo, hi, x):\n", " return x # identity function\n", "\n", " def vjp_fwd(self, nzs_in, lo, hi, x):\n", " return self(lo, hi, x), (lo, hi) # save bounds as residuals\n", "\n", " def vjp_bwd_retval(self, res, g):\n", " lo, hi = res\n", " return (None, None, jnp.clip(g, lo, hi)) # None: zero cotangents for lo, hi\n", "\n", " def batch(self, axis_data, args, dims):\n", " lo, hi, x = batch_dims_at_front(axis_data, args, dims)\n", " lo, hi = (b.reshape(b.shape[:1] + (1,) * (x.ndim - b.ndim) + b.shape[1:])\n", " for b in (lo, hi))\n", " return clip_gradient(lo, hi, x), 0\n", "\n", "def clip_gradient(lo, hi, x):\n", " return ClipGradient(jax.typeof(lo), jax.typeof(hi), jax.typeof(x))(lo, hi, x)" ] }, { "cell_type": "markdown", "id": "04e14192", "metadata": {}, "source": [ "(Static parameters in `self.params` wouldn't work for the bounds: params\n", "are baked into the primitive instance when it's constructed, so they can't\n", "be traced values. Anything that might be dynamic data — like bounds passed\n", "as arguments to a `jit`-compiled function — should be an ordinary input.\n", "Save `self.params` for genuinely static data, like the Python function in\n", "the `fixed_point` example below.)\n", "\n", "The `batch` method, for `vmap`, applies the primitive to the batched\n", "arguments. It moves every batch dimension to the front, broadcasting\n", "unbatched arguments, and then gives the bounds singleton axes after the\n", "batch axis, so that `jnp.clip` still lines up each example's bounds with\n", "that example's `x`. Broadcasting an unbatched `x` keeps the output batched,\n", "so each example's cotangent gets clipped by that example's bounds. (More\n", "on `batch` below.)" ] }, { "cell_type": "code", "execution_count": null, "id": "dbe9c998", "metadata": {}, "outputs": [], "source": [ "import matplotlib.pyplot as plt\n", "\n", "t = jnp.linspace(0, 10, 1000)\n", "\n", "plt.plot(jnp.sin(t))\n", "plt.plot(vmap(grad(jnp.sin))(t))" ] }, { "cell_type": "code", "execution_count": null, "id": "4285c068", "metadata": {}, "outputs": [], "source": [ "def clip_sin(x):\n", " x = clip_gradient(-0.75, 0.75, x)\n", " return jnp.sin(x)\n", "\n", "plt.plot(clip_sin(t))\n", "plt.plot(vmap(grad(clip_sin))(t))" ] }, { "cell_type": "markdown", "id": "2b765f19", "metadata": {}, "source": [ "### Example: Python debugging\n", "\n", "Another application that is motivated by development workflow rather than\n", "numerics is to set a `pdb` debugger trace in the backward pass of\n", "reverse-mode autodiff.\n", "\n", "When trying to track down the source of a `nan` runtime error, or just\n", "examine carefully the cotangent (gradient) values being propagated, it can be\n", "useful to insert a debugger at a point in the backward pass that corresponds\n", "to a specific point in the primal computation:" ] }, { "cell_type": "code", "execution_count": null, "id": "da77ab1b", "metadata": {}, "outputs": [], "source": [ "import pdb\n", "\n", "class Debug(HiPrim):\n", " def __init__(self, x_aval):\n", " self.in_avals = (x_aval,)\n", " self.out_aval = x_aval\n", " self.params = {}\n", " super().__init__()\n", "\n", " def expand(self, x):\n", " return x # acts like identity\n", "\n", " def vjp_fwd(self, nzs_in, x):\n", " return self(x), x\n", "\n", " def vjp_bwd_retval(self, x, g):\n", " pdb.set_trace()\n", " return (g,)\n", "\n", "def debug(x):\n", " return Debug(jax.typeof(x))(x)\n", "\n", "def foo(x):\n", " y = x ** 2\n", " y = debug(y) # insert pdb in corresponding backward pass step\n", " return jnp.sin(y)" ] }, { "cell_type": "markdown", "id": "dde6639a", "metadata": {}, "source": [ "```python\n", "jax.grad(foo)(3.)\n", "\n", "> (12)vjp_bwd_retval()\n", "-> return (g,)\n", "(Pdb) p x\n", "Array(9., dtype=float32)\n", "(Pdb) p g\n", "Array(-0.91113025, dtype=float32)\n", "(Pdb) q\n", "```\n", "\n", "### Example: Implicit function differentiation of iterative implementations\n", "\n", "This example gets pretty deep in the mathematical weeds!\n", "\n", "Another application for custom VJP rules is reverse-mode differentiation of\n", "functions that are JAX-transformable (by `jit`, `vmap`, ...) but not\n", "efficiently JAX-differentiable for some reason, perhaps because they involve\n", "`lax.while_loop`. (It's not possible to produce an XLA HLO program that\n", "efficiently computes the reverse-mode derivative of an XLA HLO While loop\n", "because that would require a program with unbounded memory use, which isn't\n", "possible to express in XLA HLO, at least without side-effecting interactions\n", "through infeed/outfeed.)\n", "\n", "For example, consider this `fixed_point` routine which computes a fixed\n", "point by iteratively applying a function in a `while_loop`:" ] }, { "cell_type": "code", "execution_count": null, "id": "8ef61fcd", "metadata": {}, "outputs": [], "source": [ "from jax.lax import while_loop\n", "\n", "def fixed_point(f, a, x_guess):\n", " def cond_fun(carry):\n", " x_prev, x = carry\n", " return jnp.abs(x_prev - x) > 1e-6\n", "\n", " def body_fun(carry):\n", " _, x = carry\n", " return x, f(a, x)\n", "\n", " _, x_star = while_loop(cond_fun, body_fun, (x_guess, f(a, x_guess)))\n", " return x_star" ] }, { "cell_type": "markdown", "id": "5aad4f1b", "metadata": {}, "source": [ "This is an iterative procedure for numerically solving the equation $x =\n", "f(a, x)$ for $x$, by iterating $x_{t+1} = f(a, x_t)$ until $x_{t+1}$ is\n", "sufficiently close to $x_t$. The result $x^*$ depends on the parameters $a$,\n", "and so we can think of there being a function $a \\mapsto x^*(a)$ that is\n", "implicitly defined by equation $x = f(a, x)$.\n", "\n", "We can use `fixed_point` to run iterative procedures to convergence, for\n", "example running Newton's method to calculate square roots while only\n", "executing adds, multiplies, and divides:" ] }, { "cell_type": "code", "execution_count": null, "id": "15fa2edd", "metadata": {}, "outputs": [], "source": [ "def newton_sqrt(a):\n", " update = lambda a, x: 0.5 * (x + a / x)\n", " return fixed_point(update, a, a)" ] }, { "cell_type": "code", "execution_count": null, "id": "61ba4e14", "metadata": {}, "outputs": [], "source": [ "print(newton_sqrt(2.))" ] }, { "cell_type": "markdown", "id": "d168f26d", "metadata": {}, "source": [ "We can `vmap` or `jit` the function as well:" ] }, { "cell_type": "code", "execution_count": null, "id": "3db87826", "metadata": {}, "outputs": [], "source": [ "print(jit(vmap(newton_sqrt))(jnp.array([1., 2., 3., 4.])))" ] }, { "cell_type": "markdown", "id": "1d2df7ff", "metadata": {}, "source": [ "We can't apply reverse-mode automatic differentiation because of the\n", "`while_loop`, but it turns out we wouldn't want to anyway: instead of\n", "differentiating through the implementation of `fixed_point` and all its\n", "iterations, we can exploit the mathematical structure to do something that is\n", "much more memory-efficient (and FLOP-efficient in this case, too!). We can\n", "instead use the implicit function theorem [Prop A.25 of Bertsekas's Nonlinear\n", "Programming, 2nd ed.], which guarantees (under some conditions) the existence\n", "of the mathematical objects we're about to use. In essence, we linearize at\n", "the solution and solve those linear equations iteratively to compute the\n", "derivatives we want.\n", "\n", "Consider again the equation $x = f(a, x)$ and the function $x^*$. We want to\n", "evaluate vector-Jacobian products like $v^\\mathsf{T} \\mapsto v^\\mathsf{T}\n", "\\partial x^*(a_0)$.\n", "\n", "At least in an open neighborhood around the point $a_0$ at which we want to\n", "differentiate, let's assume that the equation $x^*(a) = f(a, x^*(a))$ holds\n", "for all $a$. Since the two sides are equal as functions of $a$, their\n", "derivatives must be equal as well, so let's differentiate both sides:\n", "\n", "$\\qquad \\partial x^*(a) = \\partial_0 f(a, x^*(a)) + \\partial_1 f(a, x^*(a))\n", "\\partial x^*(a)$.\n", "\n", "Setting $A = \\partial_1 f(a_0, x^*(a_0))$ and $B = \\partial_0 f(a_0,\n", "x^*(a_0))$, we can write the quantity we're after more simply as\n", "\n", "$\\qquad \\partial x^*(a_0) = B + A \\partial x^*(a_0)$,\n", "\n", "or, by rearranging,\n", "\n", "$\\qquad \\partial x^*(a_0) = (I - A)^{-1} B$.\n", "\n", "That means we can evaluate vector-Jacobian products like\n", "\n", "$\\qquad v^\\mathsf{T} \\partial x^*(a_0) = v^\\mathsf{T} (I - A)^{-1} B =\n", "w^\\mathsf{T} B$,\n", "\n", "where $w^\\mathsf{T} = v^\\mathsf{T} (I - A)^{-1}$, or equivalently\n", "$w^\\mathsf{T} = v^\\mathsf{T} + w^\\mathsf{T} A$, or equivalently\n", "$w^\\mathsf{T}$ is the fixed point of the map $u^\\mathsf{T} \\mapsto\n", "v^\\mathsf{T} + u^\\mathsf{T} A$. That last characterization gives us a way to\n", "write the VJP for `fixed_point` in terms of a call to `fixed_point`!\n", "Moreover, after expanding $A$ and $B$ back out, we can see we need only to\n", "evaluate VJPs of $f$ at $(a_0, x^*(a_0))$.\n", "\n", "Here's the upshot. The function argument `f` isn't differentiated, so it goes\n", "in `params` (functions are hashable), while `a` and `x_guess` are ordinary\n", "traced inputs:" ] }, { "cell_type": "code", "execution_count": null, "id": "9b74a552", "metadata": {}, "outputs": [], "source": [ "from functools import partial\n", "from jax import vjp\n", "\n", "class FixedPoint(HiPrim):\n", " def __init__(self, a_aval, x_aval, *, f):\n", " self.in_avals = (a_aval, x_aval)\n", " self.out_aval = x_aval\n", " self.params = dict(f=f)\n", " super().__init__()\n", "\n", " def expand(self, a, x_guess):\n", " def cond_fun(carry):\n", " x_prev, x = carry\n", " return jnp.abs(x_prev - x) > 1e-6\n", "\n", " def body_fun(carry):\n", " _, x = carry\n", " return x, self.f(a, x)\n", "\n", " _, x_star = while_loop(cond_fun, body_fun, (x_guess, self.f(a, x_guess)))\n", " return x_star\n", "\n", " def vjp_fwd(self, nzs_in, a, x_guess):\n", " x_star = self(a, x_guess)\n", " return x_star, (a, x_star)\n", "\n", " def vjp_bwd_retval(self, res, x_star_bar):\n", " a, x_star = res\n", " _, vjp_a = vjp(lambda a: self.f(a, x_star), a)\n", " a_bar, = vjp_a(fixed_point(partial(rev_iter, self.f),\n", " (a, x_star, x_star_bar),\n", " x_star_bar))\n", " return a_bar, jnp.zeros_like(x_star)\n", "\n", "def rev_iter(f, packed, u):\n", " a, x_star, x_star_bar = packed\n", " _, vjp_x = vjp(lambda x: f(a, x), x_star)\n", " return x_star_bar + vjp_x(u)[0]\n", "\n", "def fixed_point(f, a, x_guess):\n", " a_aval = jax.tree.map(jax.typeof, a)\n", " x_aval = jax.tree.map(jax.typeof, x_guess)\n", " return FixedPoint(a_aval, x_aval, f=f)(a, x_guess)" ] }, { "cell_type": "code", "execution_count": null, "id": "b99c9762", "metadata": {}, "outputs": [], "source": [ "print(newton_sqrt(2.))" ] }, { "cell_type": "code", "execution_count": null, "id": "ba7a4dfd", "metadata": {}, "outputs": [], "source": [ "print(grad(newton_sqrt)(2.))\n", "print(grad(grad(newton_sqrt))(2.))" ] }, { "cell_type": "markdown", "id": "2aca2528", "metadata": {}, "source": [ "We can check our answers by differentiating `jnp.sqrt`, which uses a totally\n", "different implementation:" ] }, { "cell_type": "code", "execution_count": null, "id": "020d64ec", "metadata": {}, "outputs": [], "source": [ "print(grad(jnp.sqrt)(2.))\n", "print(grad(grad(jnp.sqrt))(2.))" ] }, { "cell_type": "markdown", "id": "7d85284a", "metadata": {}, "source": [ "Notice that the backward rule calls `fixed_point` again, on the linear\n", "problem, and that the parameter `a` passed there is a pytree of arrays: the\n", "`in_avals` entries can themselves be pytrees of types, as discussed below.\n", "\n", "A limitation to this approach is that the argument `f` can't close over any\n", "values involved in differentiation, since it's a static parameter of the\n", "primitive. That is, you might notice that we kept the parameter `a` explicit\n", "in the argument list of `fixed_point`. For this use case, consider using the\n", "low-level primitive `lax.custom_root`, which allows for derivatives in\n", "closed-over variables with custom root-finding functions.\n", "\n", "## Basic usage of `HiPrim`\n", "\n", "### Anatomy of a hijax primitive\n", "\n", "A hijax primitive is a subclass of `HiPrim`. Its `__init__` must set\n", "three attributes before calling `super().__init__()`:\n", "\n", "* `in_avals`, a tuple with one entry per positional argument, giving each\n", " argument's type (each entry can also be a pytree of types);\n", "* `out_aval`, the output type (also possibly a pytree of types);\n", "* `params`, a dict of hashable static parameters, which are made available as\n", " attributes on the instance (e.g. `self.f` for `params = dict(f=f)`).\n", " Params are static: they're baked into the primitive instance, so they\n", " can't be traced values, and dynamic data must instead be an input.\n", "\n", "Since the types are fixed at construction time, the usual idiom is a wrapper\n", "function that builds the primitive instance from the types of the arguments,\n", "using `jax.typeof`, and immediately applies it:" ] }, { "cell_type": "code", "execution_count": null, "id": "de380b6a", "metadata": {}, "outputs": [], "source": [ "class Square(HiPrim):\n", " def __init__(self, x_aval):\n", " self.in_avals = (x_aval,)\n", " self.out_aval = x_aval\n", " self.params = {}\n", " super().__init__()\n", "\n", " def expand(self, x):\n", " return x * x\n", "\n", "def square(x):\n", " return Square(jax.typeof(x))(x)\n", "\n", "print(square(3.))" ] }, { "cell_type": "markdown", "id": "f1694bf7", "metadata": {}, "source": [ "The instance is callable, and calling it is what binds the primitive: in an\n", "eager context it evaluates, and under a `jit` trace it records itself into\n", "the jaxpr as a single equation. The types of the actual arguments are checked\n", "against `in_avals`.\n", "\n", "The only required method is `expand`, which gives the implementation as a\n", "JAX-traceable Python function of the (non-static) arguments. Everything else\n", "is optional, and is only needed if you use the corresponding transformation:\n", "\n", "| method(s) | transformation |\n", "|---|---|\n", "| `expand` | evaluation and lowering (`jit`) |\n", "| `vjp_fwd` and `vjp_bwd_retval` (or `vjp_bwd`) | reverse-mode autodiff (`grad`, `vjp`) |\n", "| `jvp` | forward-mode autodiff (`jvp`) |\n", "| `lin` and `linearized` | `jax.linearize` |\n", "| `batch` | `vmap` |\n", "| `transpose` | transposition, for primitives linear in some inputs |\n", "\n", "If you apply a transformation without having defined the corresponding\n", "method, you get a `NotImplementedError` telling you what to implement. For a\n", "primitive with a `jvp` rule, the reverse-mode and linearize methods can also\n", "be derived automatically — see the helpers below.\n", "\n", "### Custom VJPs with `vjp_fwd` and `vjp_bwd_retval`\n", "\n", "The pair `vjp_fwd`/`vjp_bwd_retval` works just like the `f_fwd`/`f_bwd` pair\n", "of `jax.custom_vjp`, in Haskell-like type signatures:\n", "\n", "```haskell\n", "vjp_fwd :: (NonZeros, a) -> (b, c)\n", "vjp_bwd_retval :: (c, CT b) -> CT a\n", "```\n", "\n", "The function `vjp_fwd` describes the forward pass: it takes the primal\n", "inputs and returns a pair of the primal output and any \"residual\" data to be\n", "stored for use by the backward pass. (Its extra first argument `nzs_in` is a\n", "tuple of booleans indicating which inputs are being differentiated; you can\n", "ignore it, or use it to avoid saving residuals that won't be needed.) The\n", "primal output should usually be computed by calling `self(...)`, i.e. by\n", "binding the primitive itself; that way the custom rules also apply under\n", "higher-order differentiation.\n", "\n", "The function `vjp_bwd_retval` describes the backward pass: it takes the\n", "residuals and the cotangent of the output, and returns a tuple of cotangents\n", "with one entry per primal input." ] }, { "cell_type": "code", "execution_count": null, "id": "dd1a397c", "metadata": {}, "outputs": [], "source": [ "class Mul(HiPrim):\n", " def __init__(self, x_aval, y_aval):\n", " self.in_avals = (x_aval, y_aval)\n", " self.out_aval = x_aval\n", " self.params = {}\n", " super().__init__()\n", "\n", " def expand(self, x, y):\n", " return x * y\n", "\n", " def vjp_fwd(self, nzs_in, x, y):\n", " return self(x, y), (x, y)\n", "\n", " def vjp_bwd_retval(self, res, g):\n", " x, y = res\n", " return (g * y, x * g)\n", "\n", "def mul(x, y):\n", " return Mul(jax.typeof(x), jax.typeof(y))(x, y)\n", "\n", "print(grad(mul)(2., 3.))\n", "print(grad(mul, 1)(2., 3.))" ] }, { "cell_type": "markdown", "id": "f3dd8841", "metadata": {}, "source": [ "### Custom JVPs with `jvp`\n", "\n", "The `jvp` method defines forward-mode differentiation. It takes a tuple of\n", "primal inputs and a tuple of tangent inputs, and returns a pair of the primal\n", "output and the tangent output. (The input tangents can be symbolic zeros in\n", "some cases; see the symbolic zeros section below.)" ] }, { "cell_type": "code", "execution_count": null, "id": "da467fe5", "metadata": {}, "outputs": [], "source": [ "class Sin(HiPrim):\n", " def __init__(self, x_aval):\n", " self.in_avals = (x_aval,)\n", " self.out_aval = x_aval\n", " self.params = {}\n", " super().__init__()\n", "\n", " def expand(self, x):\n", " return jnp.sin(x)\n", "\n", " def jvp(self, primals, tangents):\n", " (x,), (x_dot,) = primals, tangents\n", " return self(x), jnp.cos(x) * x_dot\n", "\n", "def sin(x):\n", " return Sin(jax.typeof(x))(x)\n", "\n", "y, y_dot = jvp(sin, (3.,), (1.,))\n", "print(y)\n", "print(y_dot)" ] }, { "cell_type": "markdown", "id": "079f187f", "metadata": {}, "source": [ "One difference from `jax.custom_jvp`: by default, JAX does not automatically\n", "derive reverse-mode differentiation from a hijax primitive's `jvp` rule, so\n", "applying `grad` to `sin` as defined above raises a `NotImplementedError`\n", "asking for `vjp_fwd`. You can define both sets of rules on the same primitive\n", "(as in the TL;DR example above); unlike with `jax.custom_jvp` and\n", "`jax.custom_vjp`, you never have to choose between them.\n", "\n", "But you can also opt into `jax.custom_jvp`-style behavior, where everything\n", "is derived from the JVP rule, by assigning helpers as class attributes. Each\n", "helper is a pair that unpacks to the two methods it defines: `lin,\n", "linearized = linearize_from_jvp` derives linearization support by partially\n", "evaluating `jvp`, and `vjp_fwd, vjp_bwd_retval = vjp_from_lin` derives\n", "reverse mode from that linearization, storing its residuals and transposing\n", "`linearized` on the backward pass (just as `jax.custom_jvp` does):" ] }, { "cell_type": "code", "execution_count": null, "id": "05d39a1f", "metadata": {}, "outputs": [], "source": [ "from jax.experimental.hijax import linearize_from_jvp, vjp_from_lin\n", "\n", "class SinAD(Sin):\n", " lin, linearized = linearize_from_jvp\n", " vjp_fwd, vjp_bwd_retval = vjp_from_lin\n", "\n", "def sin(x):\n", " return SinAD(jax.typeof(x))(x)\n", "\n", "print(grad(sin)(3.))\n", "y, sin_lin = jax.linearize(sin, 3.)\n", "print(sin_lin(1.))\n", "print(grad(grad(sin))(3.))" ] }, { "cell_type": "markdown", "id": "bf827238", "metadata": {}, "source": [ "As with `jax.custom_jvp`, for the derived transposition to work, the JVP\n", "rule's output tangents must be linear as a function of the input tangents.\n", "\n", "A second pair, `vjp_fwd, vjp_bwd_retval = vjp_from_jvp`, instead derives\n", "reverse mode directly from the `jvp` rule: it stores the primal inputs\n", "rather than the linearization's residuals, and linearizes and transposes\n", "`jvp` on the backward pass — recomputing the rule's intermediates there,\n", "remat-style. And `jvp = jvp_from_lin` goes the other way, deriving forward mode\n", "from handwritten `lin`/`linearized` rules — a single function rather than a\n", "pair, since it defines a single method. (Deriving in a circle, with `jvp =\n", "jvp_from_lin` and `lin, linearized = linearize_from_jvp` on the same\n", "primitive, is an error.)\n", "\n", "Both kinds of rules are also what make higher-order differentiation work:\n", "`grad`-of-`grad` composes the VJP rules, while e.g. `jax.hessian`, which is\n", "forward-over-reverse, needs the `jvp` rule as well. As with `jax.custom_jvp`\n", "and `jax.custom_vjp`, a rule applies at higher orders of differentiation only\n", "if it computes its primal output by binding the primitive itself, i.e. by\n", "calling `self(...)` rather than inlining the implementation. (This represents\n", "a kind of fundamental tradeoff, where we can't make use of intermediate\n", "values from the evaluation of `expand` in our rule *and also* have the rule\n", "apply in all orders of higher-order differentiation.)\n", "\n", "### Hijax primitives in jaxprs\n", "\n", "Because a hijax primitive is a real primitive, it appears as a single\n", "equation in jaxprs:" ] }, { "cell_type": "code", "execution_count": null, "id": "68faef22", "metadata": {}, "outputs": [], "source": [ "jit(mul).trace(2., 3.).jaxpr" ] }, { "cell_type": "markdown", "id": "cb8251b8", "metadata": {}, "source": [ "That's true under `jit` too. It's only at lowering time that `expand` is\n", "traced and inlined. In an eager context, each call to the primitive calls\n", "`expand` again (so, like any JAX function, it's best to keep `expand` free of\n", "side effects, though harmless ones like `print` can be instructive):" ] }, { "cell_type": "code", "execution_count": null, "id": "d79c5d8d", "metadata": {}, "outputs": [], "source": [ "class Noisy(HiPrim):\n", " def __init__(self, x_aval):\n", " self.in_avals = (x_aval,)\n", " self.out_aval = x_aval\n", " self.params = {}\n", " super().__init__()\n", "\n", " def expand(self, x):\n", " print('called expand!')\n", " return jnp.sin(x)\n", "\n", "def noisy(x):\n", " return Noisy(jax.typeof(x))(x)\n", "\n", "print(noisy(3.))" ] }, { "cell_type": "code", "execution_count": null, "id": "4151cc59", "metadata": {}, "outputs": [], "source": [ "print(jit(noisy)(3.))\n", "print(jit(noisy)(3.)) # tracing is cached: no more 'called expand!'" ] }, { "cell_type": "markdown", "id": "b089f83c", "metadata": {}, "source": [ "You can also use Python control flow in `expand` and in the derivative\n", "rules, as long as the primitive is used eagerly (outside of `jit`), since the\n", "rules then see concrete values:" ] }, { "cell_type": "code", "execution_count": null, "id": "19fcac9d", "metadata": {}, "outputs": [], "source": [ "class G(HiPrim):\n", " def __init__(self, x_aval):\n", " self.in_avals = (x_aval,)\n", " self.out_aval = x_aval\n", " self.params = {}\n", " super().__init__()\n", "\n", " def expand(self, x):\n", " if x > 0:\n", " return jnp.sin(x)\n", " else:\n", " return jnp.cos(x)\n", "\n", " def vjp_fwd(self, nzs_in, x):\n", " return self(x), x\n", "\n", " def vjp_bwd_retval(self, x, g):\n", " if x > 0:\n", " return (2 * g,)\n", " else:\n", " return (3 * g,)\n", "\n", "def g(x):\n", " return G(jax.typeof(x))(x)\n", "\n", "print(grad(g)(1.))\n", "print(grad(g)(-1.))" ] }, { "cell_type": "markdown", "id": "1bbb1f90", "metadata": {}, "source": [ "### `vmap` with `batch`\n", "\n", "For `vmap` support, define the `batch` method. It takes the axis metadata,\n", "the batched argument values, and their batch dimensions (`None` for\n", "unbatched arguments), and returns the batched output paired with its batch\n", "dimension, computed however you like in ordinary JAX operations. Often the\n", "simplest way is to apply the primitive itself to the batched arguments, after\n", "moving their batch dimensions to the front with a helper like\n", "`batch_dims_at_front` from the gradient clipping example. That works when the\n", "primitive's rules also handle arrays with an extra leading axis, as\n", "elementwise ones do, and it keeps the batched computation a single\n", "application of the primitive, with the primitive's own rules for autodiff:" ] }, { "cell_type": "code", "execution_count": null, "id": "074bdc78", "metadata": {}, "outputs": [], "source": [ "class MulV(Mul):\n", " def batch(self, axis_data, args, dims):\n", " x, y = batch_dims_at_front(axis_data, args, dims)\n", " return mul(x, y), 0\n", "\n", "def mul(x, y):\n", " return MulV(jax.typeof(x), jax.typeof(y))(x, y)\n", "\n", "x = jnp.arange(3.)\n", "y = jnp.arange(3.) + 1.\n", "print(vmap(mul)(x, y))\n", "print(vmap(mul, in_axes=(0, None))(x, 2.))\n", "print(vmap(grad(mul))(x, y))" ] }, { "cell_type": "markdown", "id": "d1861c8d", "metadata": {}, "source": [ "`batch` can also compute the batched result some other way, for example by\n", "calling a hand-written kernel's dedicated batched variant.\n", "\n", "## More features and details\n", "\n", "### Working with `list` / `tuple` / `dict` containers (and other pytrees)\n", "\n", "The entries of `in_avals`, and `out_aval` itself, can be\n", "[pytrees](https://docs.jax.dev/en/latest/pytrees.html) of types, so arguments\n", "and outputs can be pytrees of arrays. Here's a contrived example:" ] }, { "cell_type": "code", "execution_count": null, "id": "99c45447", "metadata": {}, "outputs": [], "source": [ "from collections import namedtuple\n", "Point = namedtuple(\"Point\", [\"x\", \"y\"])\n", "\n", "from jax.experimental.hijax import Zero, instantiate_zeros\n", "\n", "class FPt(HiPrim):\n", " def __init__(self, pt_aval):\n", " self.in_avals = (pt_aval,)\n", " self.out_aval = {'a': pt_aval.x, 'b': (pt_aval.x, pt_aval.y)}\n", " self.params = {}\n", " super().__init__()\n", "\n", " def expand(self, pt):\n", " return {'a': pt.x ** 2, 'b': (jnp.sin(pt.x), jnp.cos(pt.y))}\n", "\n", " def vjp_fwd(self, nzs_in, pt):\n", " return self(pt), pt\n", "\n", " def vjp_bwd_retval(self, pt, g):\n", " g = jax.tree.map(instantiate_zeros, g,\n", " is_leaf=lambda x: isinstance(x, Zero))\n", " a_bar, (b0_bar, b1_bar) = g['a'], g['b']\n", " x_bar = 2 * pt.x * a_bar + jnp.cos(pt.x) * b0_bar\n", " y_bar = -jnp.sin(pt.y) * b1_bar\n", " return (Point(x_bar, y_bar),)\n", "\n", "def f(pt):\n", " return FPt(jax.tree.map(jax.typeof, pt))(pt)\n", "\n", "def fun(pt):\n", " dct = f(pt)\n", " return dct['a'] + dct['b'][0]\n", "\n", "pt = Point(1., 2.)\n", "print(f(pt))\n", "print(grad(fun)(pt))" ] }, { "cell_type": "markdown", "id": "fa1b84db", "metadata": {}, "source": [ "### Symbolic zeros\n", "\n", "The example above snuck in a new detail: the cotangents passed to the\n", "backward rule can contain *symbolic zeros*. When part of the primitive's\n", "output doesn't affect the value being differentiated (here `fun` doesn't use\n", "`dct['b'][1]`), the corresponding cotangent is not an array of zeros but an\n", "instance of the special `Zero` class, which records only the type. That's in\n", "contrast to `jax.custom_vjp`, where symbolic zeros are opt-in via\n", "`symbolic_zeros=True`.\n", "\n", "Symbolic zeros let a rule avoid doing work (or avoid errors, for\n", "non-differentiable outputs like integer values). If you don't want to handle\n", "them, instantiate them into actual zero arrays with `instantiate_zeros`, as\n", "above.\n", "\n", "On the input side, the `nzs_in` argument to `vjp_fwd` reports symbolically\n", "which inputs are being differentiated: it's a tuple of booleans, one per\n", "input, where `False` means that input's tangent is symbolically zero, so no\n", "cotangent for it will be used. A rule can use that to avoid saving unneeded\n", "residuals:" ] }, { "cell_type": "code", "execution_count": null, "id": "457ec431", "metadata": {}, "outputs": [], "source": [ "class Mul2(HiPrim):\n", " def __init__(self, x_aval, y_aval):\n", " self.in_avals = (x_aval, y_aval)\n", " self.out_aval = x_aval\n", " self.params = {}\n", " super().__init__()\n", "\n", " def expand(self, x, y):\n", " return x * y\n", "\n", " def vjp_fwd(self, nzs_in, x, y):\n", " x_nz, y_nz = nzs_in\n", " return self(x, y), (x if y_nz else None, y if x_nz else None)\n", "\n", " def vjp_bwd_retval(self, res, g):\n", " x, y = res\n", " return (g * y if y is not None else None,\n", " x * g if x is not None else None)\n", "\n", "def mul2(x, y):\n", " return Mul2(jax.typeof(x), jax.typeof(y))(x, y)\n", "\n", "print(grad(mul2, 0)(2., 3.)) # nzs_in == (True, False), saves only y\n", "print(grad(mul2, 1)(2., 3.)) # nzs_in == (False, True), saves only x" ] }, { "cell_type": "markdown", "id": "c75bf16c", "metadata": {}, "source": [ "Notice that the backward rule can return `None` for an input whose cotangent\n", "isn't needed.\n", "\n", "Symbolic zeros appear on the forward-mode side too: the tangents passed to a\n", "`jvp` rule can be `Zero`s, for example when a primitive is applied to a mix\n", "of differentiated inputs and constants. That's true whether the rule is\n", "invoked directly by `jax.jvp` or partially evaluated by\n", "`linearize_from_jvp`. As with cotangents, a rule can handle them explicitly\n", "to exploit the sparsity, or clean them up with `instantiate_zeros`.\n", "\n", "### In-place cotangent accumulation with `vjp_bwd`\n", "\n", "Instead of `vjp_bwd_retval`, which returns the input cotangents as values, a\n", "backward rule can be written as `vjp_bwd`, which receives one *accumulator*\n", "per primal input and adds each cotangent contribution in-place by calling\n", "`.accum(...)`:" ] }, { "cell_type": "code", "execution_count": null, "id": "e95226fc", "metadata": {}, "outputs": [], "source": [ "class Mul3(HiPrim):\n", " def __init__(self, x_aval, y_aval):\n", " self.in_avals = (x_aval, y_aval)\n", " self.out_aval = x_aval\n", " self.params = {}\n", " super().__init__()\n", "\n", " def expand(self, x, y):\n", " return x * y\n", "\n", " def vjp_fwd(self, nzs_in, x, y):\n", " return self(x, y), (x, y)\n", "\n", " def vjp_bwd(self, res, g, x_acc, y_acc):\n", " x, y = res\n", " x_acc.accum(g * y)\n", " y_acc.accum(x * g)\n", "\n", "def mul3(x, y):\n", " return Mul3(jax.typeof(x), jax.typeof(y))(x, y)\n", "\n", "print(grad(mul3)(2., 3.))" ] }, { "cell_type": "markdown", "id": "7766e69b", "metadata": {}, "source": [ "For inputs that aren't being differentiated, the accumulator is a `NullAccum`\n", "whose `accum` is a no-op, so it's safe to accumulate unconditionally. The\n", "accumulator form can save memory when gradients are accumulated across many\n", "contributions. Define either `vjp_bwd` or `vjp_bwd_retval`, not both.\n", "\n", "### Custom linearization with `lin` and `linearized`\n", "\n", "`jax.linearize` doesn't use the `jvp` or VJP rules; it has its own pair of\n", "methods. The `lin` method is like `vjp_fwd`: it takes `nzs_in` and the primal\n", "inputs, and returns the primal output paired with residuals. The `linearized`\n", "method is the linear map itself: it takes the residuals and the input\n", "tangents, and returns the output tangents:" ] }, { "cell_type": "code", "execution_count": null, "id": "73d191d1", "metadata": {}, "outputs": [], "source": [ "class Sin2(Sin):\n", " def lin(self, nzs_in, x):\n", " return self(x), jnp.cos(x)\n", "\n", " def linearized(self, cos_x, x_dot):\n", " return cos_x * x_dot\n", "\n", "def sin2(x):\n", " return Sin2(jax.typeof(x))(x)\n", "\n", "y, f_lin = jax.linearize(sin2, 3.)\n", "print(y)\n", "print(f_lin(1.))" ] }, { "cell_type": "markdown", "id": "c0d533be", "metadata": {}, "source": [ "(If you don't need to control the linearization itself, recall from above\n", "that a primitive with a `jvp` rule can just set `lin, linearized =\n", "linearize_from_jvp`.)\n", "\n", "### What we haven't covered\n", "\n", "Custom derivatives are only part of the hijax story. Hijax primitives can\n", "also:\n", "\n", "* introduce new types beyond arrays, by subclassing `HiType` and registering\n", " them with `register_hitype`, with the primitive's `in_avals`/`out_aval`\n", " mentioning the new types — see {ref}`hijax-types`;\n", "* define a `transpose` rule, for primitives that are linear in some inputs;\n", "* customize rematerialization via a `remat` method, and dead code\n", " elimination via a `dce` method.\n", "\n", "Those deserve documents of their own. In the meantime, `tests/hijax_test.py`\n", "is a good source of worked examples." ] } ], "metadata": { "jupytext": { "cell_metadata_filter": "-all", "formats": "ipynb,md:myst,py:light", "main_language": "python", "notebook_metadata_filter": "nosearch" }, "nosearch": true }, "nbformat": 4, "nbformat_minor": 5 }