Session 2: Further Differentiable Programming with JAX

Authors

Joe Wallwork

Sam Avis

Adeleke Bankole 

%matplotlib inline
import matplotlib.image
import matplotlib.pyplot as plt
import numpy as np
from mpltools import annotation

If you are using a Codespace then set it up now.

Learning objectives

In today’s session we will:

  • Try out the JAX differentiable programming framework.
  • Learn about forward mode and reverse mode applied to functions of several variables.
  • Experiment with JAX in a machine learning example.

Preparations

Terminology

This course introduces the concept of differentiable programming, a.k.a. automatic differentiation (AD), or algorithmic differentiation. We will use the acronym AD henceforth.

Notation

Recall: For a differentiable mathematical function \(f:\mathbb{R}\rightarrow\mathbb{R}\) with scalar input (i.e., a single value), we make use of Leibniz notation \(\frac{\mathrm{d}f}{\mathrm{d}x}\) for its derivative.

Caution with the physics notation for derivatives \(\dot{x}\). It won’t always mean what you expect! (See later.)

Similarly, for \(m\in\mathbb{N}\) dimensional, differentiable, vector-valued function \(\mathbf{f}:\mathbb{R}\rightarrow\mathbb{R}^m\) with scalar input, we have derivative notations \(\frac{\mathrm{d}\mathbf{f}}{\mathrm{d}x}\).

For a differentiable function with vector input (i.e., multiple inputs), we use partial derivative notation. For example, if \(f:\mathbb{R}^2\rightarrow\mathbb{R}\) is written as \(f=f(x,y)\) then we have the partial derivatives \(\frac{\partial f}{\partial x}\) and \(\frac{\partial f}{\partial y}\) with respect to first and second components, respectively. We use \[\nabla f=\left(\frac{\partial f}{\partial x_1},\dots,\frac{\partial f}{\partial x_m}\right)\] to denote the vector of all such partial derivatives. Similarly for vector-valued functions with multiple inputs.

Introduction to JAX

JAX is a differentiable programming framework written in Python but with its own Domain Specific Language (DSL). It uses just-in-time (JIT) compilation, which allows for efficient computation.

JAX is being used to enable differentiable modelling in the climate modelling domain, which has historically struggled with this. For example: * https://climate-analytics-lab.github.io/projects/jaxgcm/ * https://github.com/team-ocean/veros

Keep posted for updates on the ICCS mailing list about future projects we will be working on involving JAX.

Like Autograd, JAX overloads much of NumPy. To avoid confusion with standard NumPy and autograd.numpy, let’s import jax.numpy as jnp rather than np. One key difference with Autograd is that, unlike standard NumPy arrays, jnp arrays are immutable. This ensures that the chain of operations remains intact and allows for efficient computation.

Further similarities with Autograd are that JAX uses functional programming and the basic gradient computation function is called grad and can be imported from the main library directly.

import jax.numpy as jnp
from jax import grad

Vectorising functions with vmap

Another key JAX function is vmap, which takes a function with scalar inputs and outputs and vectorises it. The result is equivalent to if the scalar function were applied on each element of the input array(s) individually. However, the implementation is generally more efficient than doing so manually.

Below is a simple demonstration.

from jax import vmap

def square(x):
    return x ** 2

x = jnp.arange(5.0)
print(f"x = {x}")
print(f"x^2 = {vmap(square)(x)}")

Computing higher-order derivatives of a scalar function in JAX

Let’s jump straight in and consider again the example of computing derivatives of the hyperbolic tangent function, but this time with JAX instead of Autograd.

Recall that \(\tanh\) is defined by

\[\tanh(x)=\frac{1-e^{-2x}}{1+e^{-2x}}.\]

This can be implemented as the Python function:

def tanh(x):
    return (1.0 - jnp.exp((-2 * x))) / (1.0 + jnp.exp(-(2 * x)))

Exercise

Below is the same code for plotting \(\tanh\) over the range \([-7,7]\). It’s the same example we saw yesterday with Autograd. Add code for printing its first four derivatives over the same range using JAX.

Hint: JAX doesn’t have a direct equivalent of Autograd’s elementwise_grad function. Combine calls to JAX’s grad function with a call to vmap at the outer-most level, which automatically transforms functions into their ‘batched’ version, allowing you to compute gradients of each entry in an array.

Solution

from jax import grad, vmap

x = np.linspace(-7, 7, 700)

fig, axes = plt.subplots()
axes.plot(x, tanh(x), label=r"$\tanh(x)$")
axes.plot(x, vmap(grad(tanh))(x), label="First derivative")
axes.plot(x, vmap(grad(grad(tanh)))(x), label="Second derivative")
axes.plot(x, vmap(grad(grad(grad(tanh))))(x), label="Third derivative")
axes.plot(x, vmap(grad(grad(grad(grad(tanh)))))(x), label="Fourth derivative")
axes.set_xlabel(r"$x$")
axes.set_ylabel(r"$y$")
axes.set_xlim([-7, 7])
axes.grid()
axes.legend();
from jax import grad, vmap

x = np.linspace(-7, 7, 700)

fig, axes = plt.subplots()
axes.plot(x, tanh(x), label=r"$\tanh(x)$")
# TODO: Plot first-, second-, third-, and fourth-order derivatives, too
axes.set_xlabel(r"$x$")
axes.set_ylabel(r"$y$")
axes.set_xlim([-7, 7])
axes.grid()
axes.legend();

Recall: Forward and reverse mode for scalar functions

In session 1, given three scalar functions \(f\), \(g\), and \(h\), which may be composed as \[\ell(x)=h(g(f(x)),\] we presented forward mode as evaluating from right-to-left: \[\frac{\mathrm{d}\ell}{\mathrm{d}x}=\frac{\mathrm{d}h}{\mathrm{d}g}\left(\frac{\mathrm{d}g}{\mathrm{d}f}\frac{\mathrm{d}f}{\mathrm{d}x}\right)\] and reverse mode as evaluating from left-to-right: \[\frac{\mathrm{d}\ell}{\mathrm{d}x}=\left(\frac{\mathrm{d}h}{\mathrm{d}g}\frac{\mathrm{d}g}{\mathrm{d}f}\right)\frac{\mathrm{d}f}{\mathrm{d}x}.\]

Forward mode for vector functions

Suppose we have a function mapping between vectors, \(\mathbf{f}:\mathbb{R}^m\rightarrow\mathbb{R}^n\) and a point \(\mathbf{x}\in\mathbb{R}^m\) at which we seek to evaluate its derivative.

Note We can interpret such a function as having one or more scalar inputs and one or more scalar outputs.

For a general definition of forward mode, we need to consider a seed vector, \(\dot{\mathbf{x}}\in\mathbb{R}^m\). Forward mode allows us to compute the action (matrix-vector product) \[\text{JVP}(\mathbf{f},\mathbf{x},\dot{\mathbf{x}}):=\nabla\mathbf{f}(\mathbf{x})\,\dot{\mathbf{x}}.\] Here \(\nabla\mathbf{f}\) is referred to as the Jacobian for the map, so the above is known as a Jacobian-vector product (JVP). You might also hear the related term tangent linear model (TLM).

Think of the seed vector as the direction in which we want to compute the derivative. For example, using \(\dot{\mathbf{x}} = (1, 0, 0, ...)\) would give the derivative with respect to the first scalar input. In practice, the seed vector is often a derivative of some upstream code from outside of the part of the program being differentiated. That is, the upstream code is passive, whereas the part we are interested in is active, as far as AD is concerned.

Note The computation is matrix-free. We don’t actually need to assemble the Jacobian when we compute this product.

Question

In the scalar case we assumed a value for the seed. What was it?

\(f\) & \(g\) example

Consider two functions acting on real numbers: \[f(x_1,x_2)=x_1x_2\] and \[g(y)=(\sin(y),\cos(y)).\] Here \(f:\mathbb{R}^2\rightarrow\mathbb{R}\) takes two inputs and returns a single output, while \(g:\mathbb{R}\rightarrow\mathbb{R}^2\) takes a single input and returns two outputs.

Exercise

Convince yourself that it is well defined for these functions may be composed in either order. (Although they won’t necessarily give the same value!)

Solution

The image of \(f\) is the set of all real numbers, so its image is the same as the domain of \(g\) (i.e., \(\text{im}(f)=\mathbb{R}=\text{dom}(g)\)).

The image of \(g\) is \([-1,1]^2=[-1,1]\times[-1,1]\) because \(\sin\) and \(\cos\) give values between -1 and 1. Since this is a subset of \(\mathbb{R}^2\), the image of \(g\) is a subset of the domain of \(f\) (i.e., \(\text{im}(g)\subset\text{dom}(f)\)).

\(f\) & \(g\) example: Directed Acyclic Graph

We can visualise the functions in terms of DAGs.

Recalling that \[f(x_1,x_2)=x_1x_2\] and \[g(y)=(\sin(y),\cos(y))\, ,\] we have

Figure 1: Directed Acyclic Graph (DAG) for the \(f\) function in the \(f\) & \(g\) example. Generated using tikZ and \(\LaTeX\).

Figure 2: Directed Acyclic Graph (DAG) for the \(g\) function in the \(f\) & \(g\) example. Generated using tikZ and \(\LaTeX\).

\(f\) & \(g\) example: Jacobian vector product

Let’s implement this example in JAX. Recall that we need to make use of JAX arrays.

def f(x):
    return jnp.asarray([x[0] * x[1]])
def g(y):
    return jnp.asarray([jnp.sin(y), jnp.cos(y)])
from jax import jvp

Exercise 1

Take a look at the helptext for jvp to determine its signature. (Type help(jvp).) Use this to compute both the function evaluation value y for \(y=f(x)\) with \(x=(2,\pi)\) and the partial derivative dfdx1 for \(\frac{\partial f}{\partial x_1}\) for the same \(x\).

Solution

x = jnp.array([2.0, np.pi])
xdot = jnp.array([1.0, 0.0])
y, dfdx1 = jvp(f, (x,), (xdot,))


Exercise 2

Do the same thing for the partial derivative \(\frac{\partial f}{\partial x_2}\) for the same \(x\).

Solution

x = jnp.array([2.0, np.pi])
xdot = jnp.array([0.0, 1.0])
y, dfdx2 = jvp(f, (x,), (xdot,))


Exercise 3

Assemble the full Jacobian.

Solution

dfdx = jnp.hstack((dfdx1, dfdx2))
x = jnp.array([2.0, np.pi])
xdot = jnp.array([1.0, 0.0])
# TODO 1: y, dfdx1 = ???
print(f"f(x) = {y}")
print(f"dfdx1 = {dfdx1}")
xdot = jnp.array([0.0, 1.0])
# TODO 2: y, Jfx2 = ???
print(f"dfdx2 = {dfdx2}")
# TODO 3: dfdx = ???
print(f"∇f(x) = {dfdx}")

We can also compute the Jacobian in one shot using the jacfwd function.

from jax import jacfwd
x = jnp.array([2.0, np.pi])
print(f"∇f(x) = {jacfwd(f)(x)}")

Chaining forward mode derivatives

But how does this correspond to the case with three functions?

Suppose in addition to \(\mathbf{f}:\mathbb{R}^m\rightarrow\mathbb{R}^n\) we have \(\mathbf{g}:\mathbb{R}^n\rightarrow\mathbb{R}^k\), and their composition \(\boldsymbol{h}=\mathbf{f}\circ\mathbf{g}\), where the notation here implies \[ (\mathbf{f}\circ\mathbf{g})(\mathbf{x})=\mathbf{g}(\mathbf{f}(\mathbf{x})) \] for \(\mathbf{x}\in\mathbb{R}^m\). Then by applying the chain rule \[ \begin{align} \mathrm{JVP}(\boldsymbol{\ell},\mathbf{x},\dot{\mathbf{x}}) &=\mathrm{JVP}(\mathbf{f}\circ\mathbf{g},\,\mathbf{x},\,\dot{\mathbf{x}})\\ &=\nabla(\mathbf{f}\circ\mathbf{g})(\mathbf{x})\,\dot{\mathbf{x}}\\ &=\nabla\mathbf{g}(\mathbf{f}(\mathbf{x}))\,\left(\nabla\mathbf{f}(\mathbf{x})\,\dot{\mathbf{x}}\right)\\ &=\mathrm{JVP}(\mathbf{g},\,\mathbf{f}(\mathbf{x}),\,\mathrm{JVP}(\mathbf{f},\mathbf{x},\dot{\mathbf{x}})) \end{align} \]

So if we can compute a function and evaluate its forward mode derivative at the same time then we can compute such compositions and their derivatives efficiently, too.

\(f\) & \(g\) example: Seed vectors

Let’s revisit the DAG interpretation and consider how the derivatives work.

Figure 3: Directed Acyclic Graph (DAG) for the composition of the functions in the \(f\) & \(g\) example.

Exercise 1

Revisit the example above but now compute the full gradient of the composition \(h=f\circ g\) using two JVP applications.

Solution

x = jnp.array([2.0, np.pi])
xdot = jnp.array([1.0, 0.0])
z, dfdx1 = jvp(h, (x,), (xdot,))
print(f"h(x) =\n{z}")
print(f"dfdx1 =\n{dfdx1}")

xdot = jnp.array([0.0, 1.0])
z, dfdx2 = jvp(h, (x,), (xdot,))
print(f"dfdx2 =\n{dfdx2}")

dfdx = jnp.hstack((dfdx1, dfdx2))
print(f"∇f(x) = {dfdx}")


Exercise 2

Verify that your answer coincides with the result of applying jacfwd.

Solution

print(f"∇f(x) = {jacfwd(h)(x)}")

should give the same answer.

def h(x):
    return g(f(x))
# TODO 1: Compute the full gradient of h using two JVP applications
# TODO 2: Verify that your answer coincides with the result of applying jacfwd

Reverse mode for vector functions

Consider the same vector function as above, \(\mathbf{f}:\mathbb{R}^m\rightarrow\mathbb{R}^n\) and a point \(\mathbf{x}\in\mathbb{R}^m\).

Given \(\mathbf{x}\in\mathbb{R}^m\) and a seed vector \(\bar{\mathbf{y}}\in\mathbb{R}^n\), reverse mode AD allows us to compute the transpose action (vector-Jacobian product) \[\text{VJP}(\mathbf{f},\mathbf{x},\bar{\mathbf{y}}):=\bar{\mathbf{y}}^T\nabla\mathbf{f}(\mathbf{x}).\]

Notes:

  • This name is not used universally. You will often find Jacobian-transpose vector product in the literature.
  • The dimension of the seed vector corresponds to that of the output rather than the input.
  • Again, the computation is matrix-free. We don’t actually need the Jacobian or its transpose when we compute this product.
  • Our original definition for the scalar case assumed the same seed as above.

Exercise 1

Convince yourself that the VJP is well defined.

Solution

We have \(\nabla\mathbf{f}(\mathbf{x})\in\mathbb{R}^{n\times m}\), so \(\nabla\mathbf{f}(\mathbf{x})^T\in\mathbb{R}^{m\times n}\). Since \(\bar{\mathbf{y}}\in\mathbb{R}^n\), the dimensions are appropriate to take the VJP.


Optional exercise 2

Repeat the above exercise to convice yourself that the definition coincides with the scalar definition, i.e.

\[ \mathrm{VJP}(\mathbf{h},\mathbf{x},\bar{\mathbf{y}}) =\mathrm{VJP}(\mathbf{f},\,\mathbf{x},\,\mathrm{VJP}(\mathbf{g},\,\mathbf{f}(\mathbf{x}),\bar{\mathbf{y}})). \]

Solution

\[ \begin{align} \mathrm{VJP}(\mathbf{h},\mathbf{x},\bar{\mathbf{y}}) &=\mathrm{JVP}(\mathbf{f}\circ\mathbf{g},\,\mathbf{x},\,\bar{\mathbf{y}})\\ &=\bar{\mathbf{y}}^T\,\nabla(\mathbf{f}\circ\mathbf{g})(\mathbf{x})\\ &=\left(\nabla\mathbf{g}(\mathbf{f}(\mathbf{x}))\,\nabla\mathbf{f}(\mathbf{x})\right)^T\bar{\mathbf{y}}\\ &=\left(\nabla\mathbf{f}(\mathbf{x})^T\,\nabla\mathbf{g}(\mathbf{f}(\mathbf{x}))^T\right)\,\bar{\mathbf{y}}\\ &=\nabla\mathbf{f}(\mathbf{x})^T\,\left(\nabla\mathbf{g}(\mathbf{f}(\mathbf{x}))^T\,\bar{\mathbf{y}}\right)\\ &=\left(\bar{\mathbf{y}}^T\,\nabla\mathbf{g}(\mathbf{f}(\mathbf{x}))\right)^T\,\nabla\mathbf{f}(\mathbf{x})\\ &=\mathrm{VJP}(\mathbf{f},\,\mathbf{x},\,\mathrm{VJP}(\mathbf{g},\,\mathbf{f}(\mathbf{x}),\bar{\mathbf{y}}))) \end{align} \]

where we have used the result from before and applied the rule of transposes for matrix multiplication.

\(f\) & \(g\) example: reverse mode seed vectors

In reverse mode, gradient information propagates in the opposite direction.

Figure 4: Directed Acyclic Graph (DAG) for the composition of the functions in the \(f\) & \(g\) example. Generated using tikZ and \(\LaTeX\).

\(f\) & \(g\) example: vector-Jacobian product

Exercise 1

Take a look at the helptext for vjp to determine its signature. Use a single application to compute both the function evaluation value y for \(y=f(x)\) with \(x=(2,\pi)\) and the full gradient.

Note: The API for vjp is a little different from the API for jvp.

Solution

x = jnp.array([2.0, np.pi])
y, VJP = vjp(f, x)
ybar = jnp.array([1.0])  # Unit seed
print(f"VJP(f, x, ybar) = {VJP(ybar)}")


Exercise 2

Verify that your answer coincides with the result of applying the jacrev driver (the reverse mode counterpart to jacfwd).

Solution

from jax import jacrev

print(f"∇f(x) = {jacrev(f)(x)}")

should give the same answer.


Optional exercise 3

Compare timings for a function evaluation, the forward mode derivative computation, and the reverse mode derivative computation using the %%timeit notebook magic.

Solution

First, execute

x = jnp.array([2.0, np.pi])

in its own cell. Then

%%timeit
y = f(x)

to time the function evaluation. For the derivatives, execute

%%timeit
y, dfdx1 = jvp(f, (x,), (jnp.array([1.0, 0.0]),))
y, dfdx2 = jvp(f, (x,), (jnp.array([0.0, 1.0]),))

in one cell for forward mode and

%%timeit
y, VJP = vjp(f, x)
dfdx = VJP(jnp.array([1.0]))

in another cell for reverse mode.

You will likely find that function evaluation is much faster than derivative computation for this simple example. You will likely also find that reverse mode is slower than forward mode, despite the fact that it involves two jvp calls (as opposed to a single vjp call).

from jax import vjp
# TODO 1: Compute the full gradient of f using a single vjp application
# TODO 2: Verify that the answer coincides with that given by jacrev
# TODO 3 (optional): Compare timings between function evaluation, forward mode, and reverse mode

Forward mode vs. reverse mode

For seed vectors \(\dot{\mathbf{x}}\in\mathbb{R}^n\) and \(\bar{\mathbf{y}}\in\mathbb{R}^m\), forward mode and reverse mode compute

\[ \text{JVP}(\mathbf{f},\mathbf{x},\dot{\mathbf{x}}):=\nabla\mathbf{f}(\mathbf{x})\dot{\mathbf{x}} \quad\text{and}\quad \text{VJP}(\mathbf{f},\mathbf{x},\bar{\mathbf{y}}):=\bar{\mathbf{y}}^T\nabla\mathbf{f}(\mathbf{x}) \] respectively.

  • Forward mode is more appropriate if \(n\ll m\), i.e., \(\#inputs\ll\#outputs\).
    • e.g., sensitivity analysis or optimisation w.r.t. a small number of parameters.
  • Reverse mode is more appropriate if \(n\gg m\), i.e., \(\#inputs\gg\#outputs\).
    • e.g., ODE/PDE-constrained optimisation (objective/cost function), machine learning training (loss function), goal-oriented error estimation (quantity of interest).
  • Forward mode is computed eagerly, whereas reverse mode is done separately from the primal run.
  • Reverse mode tends to have higher memory requirements.

Note: As a rule of thumb, one reverse mode call is roughly equivalent to four evaluations of the function being differentiated.

Computing Hessians in the vector case

The Hessian is the matrix of second-order partial derivatives. For a scalar-valued function with several inputs, \(f:\mathbb{R}^n\rightarrow\mathbb{R}\), it is given by \[ \mathbf{H}(f)=\begin{bmatrix} \frac{\partial^2f}{\partial x_1^2} & \frac{\partial^2 f}{\partial x_1\partial x_2} & \dots & \frac{\partial^2 f}{\partial x_1\partial x_n}\\ \frac{\partial^2f}{\partial x_2\partial x_1} & \frac{\partial^2 f}{\partial x_2^2} & \dots & \frac{\partial^2 f}{\partial x_2\partial x_n}\\ \vdots & \vdots & \ddots & \vdots\\ \frac{\partial^2f}{\partial x_n\partial x_1} & \frac{\partial^2 f}{\partial x_n\partial x_2} & \dots & \frac{\partial^2 f}{\partial x_n^2} \end{bmatrix}. \]

For seed vectors \(\dot{\mathbf{x}}\in\mathbb{R}^n\) and \(\bar{\mathbf{y}}\in\mathbb{R}^m\), forward mode and reverse mode compute

\[ \text{JVP}(f,\mathbf{x},\dot{\mathbf{x}}):=\nabla f(\mathbf{x})\dot{\mathbf{x}} \quad\text{and}\quad \text{VJP}(f,\mathbf{x},\bar{\mathbf{y}}):=\bar{\mathbf{y}}^T\nabla f(\mathbf{x}) \] respectively.

Question

What are two ways we can we use these to compute the Hessian of \(f\)?

Solution 1

Given a seed vector \(\dot{\mathbf{x}}\), first apply forward mode to compute \(\nabla f(\mathbf{x})\dot{\mathbf{x}}\). Then apply forward mode to compute the gradient of this (i.e., apply forward mode to the forward mode derivative code). Use vector mode (preferably with compression!) to get the full Hessian.


Solution 2

Given a seed vector \(\dot{\mathbf{x}}\), first apply forward mode to compute \(\nabla f(\mathbf{x})\dot{\mathbf{x}}\). Then apply reverse mode to compute the gradient of this (i.e., apply reverse mode to the forward mode derivative code). That is, \((\nabla(\nabla f(\mathbf{x})\dot{\mathbf{x}}))^T\bar{\mathbf{y}}=\dot{\mathbf{x}}^T\nabla^T\nabla f(\mathbf{x})\bar{\mathbf{y}}\). Here the Hessian \(\mathbf{H}(f):=\nabla^T\nabla f(\mathbf{x})\) is symmetric and so the two applications give the Hessian-vector product with the seed. Use vector mode (preferably with compression!) to get the full Hessian.

\(f\) & \(g\) example: Hessian computation

Let’s apply this to the \(f\) function above. Using both approaches we should get the same answer.

from jax import jacfwd, jacrev

def hessian1(fun):
    return jacfwd(jacfwd(fun))

def hessian2(fun):
    return jacfwd(jacrev(fun))
hessian1(f)(x)
hessian2(f)(x)

Neural network example

In the previous session we trained a basic neural network to recognise handwritten numbers using autograd. Here we will use JAX to investigate this trained neural network.

Figure 5: Neural network schematic created with https://alexlenail.me/NN-SVG/index.html.


To begin with run the two blocks below to train the neural network. These do the same as last session but this time implemented using JAX.

# Functions and imports for training the handwriting classifier neural network with JAX

import array
import gzip
import os
import struct
from urllib.request import urlretrieve

import jax
import jax.numpy as jnp
from jax import grad, jit, random
from jax.scipy.special import logsumexp
import optax
import matplotlib.pyplot as plt

def load_mnist():
    """
    Download the MNIST training data and normalise it in preparation for training and testing.

    :return: number of images and arrays of normalised images and labels
    """
    print("Loading training data...")

    def download(url, filename):
        if not os.path.exists("data"):
            os.makedirs("data")
        out_file = os.path.join("data", filename)
        if not os.path.isfile(out_file):
            urlretrieve(url, out_file)

    def parse_labels(filename):
        with gzip.open(filename, "rb") as fh:
            magic, num_data = struct.unpack(">II", fh.read(8))
            return np.array(array.array("B", fh.read()), dtype=np.uint8)

    def parse_images(filename):
        with gzip.open(filename, "rb") as fh:
            magic, num_data, rows, cols = struct.unpack(">IIII", fh.read(16))
            return np.array(array.array("B", fh.read()), dtype=np.uint8).reshape(num_data, rows, cols)

    base_url = "https://storage.googleapis.com/cvdf-datasets/mnist/"
    
    for filename in [
        "train-images-idx3-ubyte.gz",
        "train-labels-idx1-ubyte.gz",
        "t10k-images-idx3-ubyte.gz",
        "t10k-labels-idx1-ubyte.gz",
    ]:
        download(base_url + filename, filename)

    train_images = parse_images("data/train-images-idx3-ubyte.gz")
    train_labels = parse_labels("data/train-labels-idx1-ubyte.gz")
    test_images = parse_images("data/t10k-images-idx3-ubyte.gz")
    test_labels = parse_labels("data/t10k-labels-idx1-ubyte.gz")

    partial_flatten = lambda x: jnp.reshape(x, (x.shape[0], -1))
    one_hot = lambda x, k: jnp.array(x[:, None] == jnp.arange(k)[None, :], dtype=int)
    train_images = partial_flatten(train_images) / 255.0
    test_images = partial_flatten(test_images) / 255.0
    train_labels = one_hot(train_labels, 10)
    test_labels = one_hot(test_labels, 10)
    N_data = train_images.shape[0]
    return N_data, train_images, train_labels, test_images, test_labels

def init_random_params(scale, layer_sizes, key=random.PRNGKey(0)):
    """
    Build a list of (weights, biases) tuples, one for each layer in the net.

    :arg scale: scaling parameter
    :arg layer_sizes: sizes of each layer in the neural network
    :kwarg rs: initial random state
    :return: array of (weights, biases) tuples representing the net
    """
    keys = random.split(key, len(layer_sizes) - 1)
    params = []
    for k, m, n in zip(keys, layer_sizes[:-1], layer_sizes[1:]):
        W = scale * random.normal(k, (m, n))  # weight matrix
        b = scale * random.normal(k, (n))     # bias vector
        params.append((W, b))
    return params

def neural_net_predict(params, inputs):
    """
    Implements a deep neural network for classification.
    
    :arg params: list of (weights, bias) tuples
    :arg inputs: (N x D) input data matrix for images
    :return: class probabilities
    """
    for W, b in params[:-1]:
        inputs = jnp.tanh(inputs @ W + b)
    W, b = params[-1]
    outputs = inputs @ W + b
    return jnp.exp(outputs - logsumexp(outputs, axis=1, keepdims=True)) # softmax

def log_posterior(params, inputs, targets, L2_reg):
    """
    Compute a log-posterior objective function comprised of prior and likelihood components

    :arg params: list of (weights, bias) tuples
    :arg inputs: input data matrix for images
    :arg targets: target data for labels
    :arg L2_reg: scalar regularisation parameter
    :return: log-posterior objective function
    """
    l2_norm = sum(jnp.sum(W**2) + jnp.sum(b**2) for W, b in params)
    log_prior = -L2_reg * l2_norm
    probs = neural_net_predict(params, inputs)
    log_lik = jnp.sum(jnp.log(probs) * targets)
    return log_prior + log_lik

def accuracy(params, inputs, targets):
    """
    Compute the accuracy of a set of predictions.

    :arg params: list of (weights, bias) tuples
    :arg inputs: input data matrix for images
    :arg targets: target data for labels
    :return: mean accuracy of predictions
    """
    target_class = jnp.argmax(targets, axis=1)
    predicted_class = jnp.argmax(neural_net_predict(params, inputs), axis=1)
    return jnp.mean(predicted_class == target_class)

# Define training objective
def objective(params, inputs, targets, L2_reg=1.0):
    return -log_posterior(params, inputs, targets, L2_reg)

def print_perf(params, iteration):
    if iteration == 0:
        print("Epoch | Train accuracy | Test accuracy")
    if iteration % num_batches == 0:
        train_acc = accuracy(params, train_images, train_labels)
        test_acc = accuracy(params, test_images, test_labels)
        print(f"  {iteration // num_batches:1}   |  {train_acc:10.5f}    |  {test_acc:10.5f}")

def get_batch(iteration):
    idx = (iteration % num_batches) * batch_size
    x = train_images[idx:idx+batch_size]
    y = train_labels[idx:idx+batch_size]
    return x, y

def show_images(
    images,
    ims_per_row=5,
    padding=5,
    digit_dimensions=(28, 28),
    cmap="gray_r",
    vmin=None,
    vmax=None,
):
    """Images should be a (N_images x pixels) matrix."""
    N_images = images.shape[0]
    N_rows = (N_images - 1) // ims_per_row + 1
    pad_value = np.min(images.ravel())
    concat_images = np.full(
        (
            (digit_dimensions[0] + padding) * N_rows + padding,
            (digit_dimensions[1] + padding) * ims_per_row + padding,
        ),
        pad_value,
    )
    for i in range(N_images):
        cur_image = np.reshape(images[i, :], digit_dimensions)
        row_ix = i // ims_per_row
        col_ix = i % ims_per_row
        row_start = padding + (padding + digit_dimensions[0]) * row_ix
        col_start = padding + (padding + digit_dimensions[1]) * col_ix
        concat_images[
            row_start : row_start + digit_dimensions[0], col_start : col_start + digit_dimensions[1]
        ] = cur_image
    
    fig, axes = plt.subplots()
    cax = axes.matshow(concat_images, cmap=cmap, vmin=vmin, vmax=vmax)
    plt.xticks(np.array([]))
    plt.yticks(np.array([]))
# Model parameters
param_scale = 0.1
layer_sizes = [784, 200, 100, 10]
init_params = init_random_params(param_scale, layer_sizes)

# Training parameters
batch_size = 256
num_epochs = 5
step_size = 0.001

## Load training data
N, train_images, train_labels, test_images, test_labels = load_mnist()
num_batches = int(np.ceil(len(train_images) / batch_size))
print(f"Training dataset size: {N}")
print(f"Testing dataset size: {len(test_images)}")

## Train the neural network    
# Set up the solver
solver = optax.adam(step_size)
opt_state = solver.init(init_params)
params = init_params

# Get gradient of the objective function using JAX.
grad_fn = grad(objective)

@jit # Compile each step of the solver for better performance
def step(params, opt_state, inputs, targets):
    grads = grad_fn(params, inputs, targets)
    updates, opt_state = solver.update(grads, opt_state)
    params = optax.apply_updates(params, updates)
    return params, opt_state

# Run the solver
for iteration in range(num_epochs * num_batches):
    inputs, targets = get_batch(iteration)
    params, opt_state = step(params, opt_state, inputs, targets)
    print_perf(params, iteration)
optimized_params = params

## Check the result
show_images(test_images[100:105])
predicted = neural_net_predict(optimized_params, test_images[100:105]).argmax(axis=1)
print(f"Predicted: {predicted}")

Neural network example: Effect of each pixel on the predicted result

Starting from a blank image we can investigate how changing each pixel will influence the number predicted by the neural network.

As an aside, we can investigate what the neural network outputs for a blank image… although it is likely to be nonsense.

blank_image = jnp.zeros(28 * 28)

# The neural network expects and returns a 2D array of probabilities for each digit for multiple images
probs = neural_net_predict(optimized_params, blank_image[np.newaxis,:])[0]

# Show the probabilities for a blank image
fig = plt.figure(figsize=(5,3))
plt.bar(range(10), probs)
plt.xticks(range(10))
plt.ylabel("Digit probabilities")
plt.ylim([0, 1])
plt.show()

Exercise 1

What function, when differentiated, will show how each pixel influences the predicted digit? Fill this in in the code below.

Hint The function should take the image as input and return the probabilities for each of the ten digits (the output of the neural network).
Solution 
def f(x):
    """
    :arg x: input image
    :return: prediction of the digit this image represents
    """
    return neural_net_predict(optimized_params, x[np.newaxis, :])[0]


Question 2

We will be calculating the derivative of all outputs (digits) with respect to all inputs (pixels), i.e. the full Jacobian. Which differentiation approach will work best, forward mode or reverse mode? And which function should be used jax.grad, jax.jvp or jax.vjp?

Hint Consider the number of inputs and outputs to the function. Generally, more inputs - reverse mode, more outputs - forward mode.
Solution

Reverse mode: Because there are more inputs (\(28 \times 28 = 784\) pixels) than outputs (\(10\)).

jax.vjp (vector-Jacobian product) because this is the driver for reverse mode. jax.jvp (Jacobian-vector product) is used for forward mode and the output is not a scalar so jax.grad cannot be used.


Exercise 3

Complete the code to get the gradients for each digit with respect to the pixels.

Hint

jax.vjp takes the function to differentiate and its input. It returns the function output (not needed) and a functon that calculates the vector-Jacobian product when given a seed vector.

The result of this is a list of gradients for each input parameter. There is only one in this case (\(x\)).
Solution 
results = []
for digit, seed_vector in enumerate(seed_vectors):
    _, vjp_fn = jax.vjp(f, blank_image)
    result = vjp_fn(seed_vector)[0]
    results.append(result)

Or, more succinctly using jax.vmap:

_, vjp_fn = jax.vjp(f, blank_image)
results = jax.vmap(lambda v: vjp_fn(v)[0])(seed_vectors)
def f(x):
    """
    :arg x: input image
    :return: probabilities of each digit for this image
    """
    # Exercise 1: Define this
    return ???

# We will evaluate the function `f` at at array of zeros, which corresponds to a blank image
blank_image = jnp.zeros(28 * 28)

# We need seed vectors for each of the ten digits, we can use the rows of an identity matrix
seed_vectors = jnp.eye(10)

# Exercise 3: Get the effect of each pixel for each of the ten digits
results = []
for digit, seed_vector in enumerate(seed_vectors):
    result = ??? # TODO
    results.append(result)

# Plot the gradient fields
fig, axes = plt.subplots(1, 10, figsize=(14,2))
for digit, result in enumerate(results):
    vmax = np.max(np.abs(np.array(result)))
    ax = axes.ravel()[digit]
    ax.imshow(result.reshape(28, 28), cmap="RdBu_r", vmax=vmax, vmin=-vmax)
    ax.set_title(digit)
    ax.set_axis_off()

Neural network example: Effect of a perturbation

Next we will look at how making a perturbation to the image would impact the likelihood of each digit. Run the code below to make a perturbation.

# Choose an image
image = test_images[0]

# Function to differentiate
def f(x):
    return neural_net_predict(optimized_params, x[np.newaxis, :])[0]

# Define a perturbation
perturbation = 1.0 * jnp.zeros((28, 28)).at[:, 13:15].set(1.0).reshape(-1) # Vertical stripe
#perturbation = 0.2 * jnp.ones_like(image)                                 # Uniform
#perturbation = 0.2 * jax.random.normal(jax.random.PRNGKey(0), (28 * 28,)) # Random noise

# Plot
fig, axes = plt.subplots(1, 2, figsize=(6,3))

cmap = axes[0].imshow(image.reshape(28,28), cmap="gray_r")
plt.colorbar(cmap, ax=axes[0])
axes[0].set_title("Original image")
axes[0].axis("off")

cmap = axes[1].imshow((image + perturbation).reshape(28,28), cmap="gray_r")
plt.colorbar(cmap, ax=axes[1])
axes[1].set_title("After perturbation")
axes[1].axis("off")

plt.show()

Exercise 4

Complete the code below to use the gradient to estimate the change in the likelihood of each digit (output of the neural network) given the perturbation.

Hint

The gradient is calculated in a single direction in the input space (the perturbation). So the differentiation should use forward mode with jax.jvp.

The function f(x) only takes a single argument so the values and tangent passed to jax.jvp should each be a list of size 1.
Solution 
probs, change_probs = jax.jvp(f, (image,), (perturbation,))
# Exercise 4: Estimate the change in likelihood of each digit as `change_probs`
probs, change_probs = ???

# Plot
fig, axes = plt.subplots(1, 2, figsize=(6,3))
axes[0].bar(range(10), probs)
axes[0].set_xticks(range(10))
axes[0].set_title("Original probabilities")
axes[0].semilogy()
axes[1].bar(range(10), change_probs / probs)
axes[1].set_xticks(range(10))
axes[1].set_title("Estimated relative change\nin probabilities")
plt.show()

Further applications

  • Sensitivity analysis.
  • Data assimilation.
  • Uncertainty quantification.
  • PDE-constrained optimisation.
  • Goal-oriented error estimation and mesh adaptation.

Mini-project

This year’s summer school includes a ‘mini-project’ session, which will be on Thursday afternoon.

If you’re interested to make use of what you’ve learnt in this course then there will be a mini-project on investigating differentiable programming applied to toy models written in JAX. We provide code for:

  1. Lorenz ’96 - an idealised model of a simplified atmosphere, demonstrating the core dynamical system.
  2. Shallow water equations - a depth-averaged approximation to the Navier-Stokes equations applied to an idealised tsunami simulation.

For a sneak preview, take a look at the miniproject subdirectory of this repo.

Summary and outlook

In today’s session we:

  • Tried out the JAX differentiable programming framework.
  • Learnt about forward and reverse modes applied to functions of several variables.
  • Experimented with JAX in a machine learning example.

European workshop on Automatic differentiation

Figure 6: Centre of Mathematical Sciences, University of Cambridge.


ICCS will be hosting the 29th European workshop on Automatic Differentiation (EuroAD) in Cambridge on the 29th-30th September 2026. It will be an informal meeting of researchers and software engineers who develop and apply automatic differentiation theory and software.

Attendees from all career stages are welcome, particularly PhD students and early career researchers and software engineers.

We welcome contributions with both theoretical and practical perspectives. A selection of possible topics include: new methods for AD, software developments and inter-comparison, AD in machine learning, and applications in science, engineering, and beyond.

Register at https://cambridge-iccs.github.io/euroad29/

References