Session 1: Introduction to Differentiable Programming with Autograd

Authors

Joe Wallwork

Sam Avis

Adeleke Bankole

%matplotlib inline

import matplotlib.image
import matplotlib.pyplot as plt
from mpltools import annotation

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

Motivation

Differentiable programming is an enabling technology. Given scientific code for computing some quantity, it allows us to generate derivatives of quantities involved in the computation without any need for deriving or hand-coding the derivative expressions involved.

Some motivating examples: * Computing the Jacobian for a nonlinear system. * Computing the gradient required for an ODE- or PDE-constrained optimisation method. * The backpropagation operation used for training machine learning models. * Computing Hessians for uncertainty quantification methods. * Solving the adjoint problems involved in data assimilation methods commonly used for weather forecasting.

Learning objectives

In today’s session we will:

Get a brief history of automatic differentiation.
Learn about the operator overloading approach.
Learn about forward mode and reverse mode applied to scalar-valued functions.
Try out the Autograd AD tool applied to some test problems.
Verify the derivatives produced by Autograd using the Taylor test.

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

For a differentiable mathematical function \(f:A\rightarrow\mathbb{R}\) with scalar input (i.e., a single value) from \(A\subseteq\mathbb{R}\), we make use of the Leibniz notation \(\frac{\mathrm{d}f}{\mathrm{d}x}\) for its derivative. We don’t use Lagrange notation \(f'(x)\) in this course. We also don’t consider functions with multiple inputs or outputs in this session.

History of forward mode

Origins of AD in 1950s.
However, it found a wider audience in the 1980s, when it became more relevant thanks to advances in both compute power and modern programming languages.
Forward mode was discovered by Wengert in 1964.

Figure 1: Header of (R. E. Wengert, 1964).

History of reverse mode

Figure 2: Schematic from (Speelpenning, 1980).

Reverse mode (a.k.a. back-propagation) was discovered by Linnainmaa in the 1970s.
The terminology ‘back-propagating error correction’ had already been introduced in 1962 by Frank Rosenblatt, but he did not know how to implement this.
Speelpenning introduced the modern formulation of reverse mode in the late 1980s.
Griewank improved the feasibility of reverse mode in 1992 by introducing checkpointing.
Back-propagation is now well-known as a key enabling technology for machine learning.

Idea

The idea of AD is to treat a model as a sequence of elementary instructions (e.g., addition, multiplication, exponentiation). Here a model could be a function or subroutine, code block, or a whole program. Elementary operations are well-understood and their derivatives are known. As such, the derivative of the whole model may be computed by composing the derivatives of each operation using the chain rule.

Recap on A-level maths: the Chain Rule

Consider two composable, differentiable (mathematical) functions, \(f\) and \(g\). By definition, this means \[h(x)=g(f(x)).\]

Then the chain rule states that the derivative of \(h\) may be computed in terms of the derivatives of \(f\) and \(g\) using the formula \[\frac{\mathrm{d}h}{\mathrm{d}x}=\frac{\mathrm{d}g}{\mathrm{d}f}\frac{\mathrm{d}f}{\mathrm{d}x}.\]

If \(h\) were to be passed to a further function then this new derivative could be found by applying the chain rule again.

Introductory example

For example, we will take the function \(h(x) = \sin(x^2)\), which is composed of \[f(x) = x^2 \quad\text{and}\quad g(y) = \sin(y).\]

The derivatives of which are \[\frac{\mathrm{d}f}{\mathrm{d}x} = 2x, \quad \frac{\mathrm{d}g}{\mathrm{d}y} = \cos(y)\]

So, using the chain rule the total derivative is \[\frac{\mathrm{d}h}{\mathrm{d}x} = \frac{\mathrm{d}g}{\mathrm{d}f} \frac{\mathrm{d}f}{\mathrm{d}x} = \cos(f(x)) \times 2x = 2x \cos(x^2).\]

This function and derivative are shown by the code below. In this example, we make use of Autograd, which allows us to automatically differentiate NumPy. As such, we first import its alias for NumPy.

import autograd.numpy as np

Autograd uses functional programming, meaning that it operates on Python functions. The first argument to most Autograd functions is itself a function.

def f(x):
    return x**2

def g(y):
    return np.sin(y)

def h(x):
    return g(f(x))

def dhdx_analytic(x):
    return 2 * x * np.cos(x**2)

x = np.linspace(0, 2*np.pi, 100)

_, axes = plt.subplots(nrows=2, figsize=(6, 6))
axes[0].plot(x, h(x))
axes[1].plot(x, dhdx_analytic(x))
axes[0].set_ylabel("h(x)")
axes[1].set_ylabel(r"$\mathrm{d}h/\mathrm{d}x$")
axes[0].grid()
axes[1].grid()
plt.show()

The aim of AD is to obtain the derivative without having to solve it analytically as we did above.

This can be done using the autograd library using the grad function:

from autograd import grad

help(grad)

dhdx_auto_fn = grad(h)

# Compute the gradient function at each value of x
dhdx_auto = [dhdx_auto_fn(xi) for xi in x]

# Plot it
fig, axes = plt.subplots(figsize=(6, 3))
axes.plot(x, dhdx_analytic(x), "-", label="analytic")
axes.plot(x, dhdx_auto, "ro", label="AD")
axes.legend(loc="best")
axes.set_ylabel(r"$\mathrm{d}h/\mathrm{d}x$")
axes.grid(True)

Note: We can evaluate dhdx_auto more compactly as follows:

from autograd import elementwise_grad as egrad

dhdx_auto = egrad(h, x)

It is likely that this is a more efficient computation.

Optional exercises:

Convince yourself that this gives the same result.
Use the %%timeit magic to compare the timings.

Verification: the Taylor test

Recall the Taylor expansion of a differentiable scalar function \(f:\mathbb{R}\rightarrow\mathbb{R}\): \[ f(x+\epsilon)=f(x)+\epsilon\frac{\mathrm{d}f}{\mathrm{d}x}(x) + \frac{1}{2!}\epsilon^2\frac{\mathrm{d}^2f}{\mathrm{d}x^2} + \dots + \frac{1}{n!}\epsilon^n\frac{\mathrm{d}^nf}{\mathrm{d}x^n}(x) + \dots, \] for some \(\epsilon\in\mathbb{R}\), i.e., \[ f(x+\epsilon)=f(x)+\epsilon\frac{\mathrm{d}f}{\mathrm{d}x}(x) + \mathcal{O}(\epsilon^2). \] This gives rise to the first-order forward difference approximation of the first derivative as follows: \[ \frac{\mathrm{d}f}{\mathrm{d}x}\approx \frac{f(x+\epsilon)-f(x)}{\epsilon}. \]

This is a rather crude approximation to a derivative, but it can do the job, given a sufficiently small \(\epsilon\).

The idea of the Taylor test is to take smaller and smaller spacing values \(\epsilon\) from a given input value and to check that the difference between the forward difference and the AD-generated result converges quadratically. That is, we verify that \[ \|f(x+\epsilon)-f(x)-\epsilon\:\texttt{dfdx}(x)\|=\mathcal{O}(\epsilon^2), \] where \(\texttt{dfdx}\) is the AD-generated result.

We return to the previous example involving the composition \(h(x)=g(f(x))\) and double-check that we are happy with the output of the AD tool for this function.

# Choose arbitrary inputs for the composition
x = 2.5

# Compute the derivative using automatic differentiation
dhdx_auto = grad(h)(x)

# Run the Taylor test over several spacing values
spacings = [1.0, 0.1, 0.01, 0.001]
errors = []
for epsilon in spacings:
    # Compute the discrepancy
    errors.append(np.linalg.norm(h(x + epsilon) - h(x) - epsilon * dhdx_auto))

# Plot the solution, demonstrating that the expected quadratic convergence is achieved
fig, axes = plt.subplots()
axes.loglog(spacings, errors, "--x")
axes.set_xlabel(r"$\epsilon$ spacing")
axes.set_ylabel(r"$\ell_2$ error")
annotation.slope_marker((1e-2, 2e-4), 2, ax=axes, invert=True)
axes.grid()

Operator overloading

What’s happening under the hood? To compute the gradient, Autograd first has to record every operation applied in the program that is relevant to the calculation. This is achieved by wrapping functions such that calling them adds them to a list of operations known as a tape or Wengert list. Abstract idea:

Number	Operation	Input	Output
1	`**2`	`x`	`y`
2	`sin`	`y`	`z`

Autograd has a table that maps these wrapped functions to their corresponding derivatives.

Operation	Derivative
`**2`	`*2`
`sin`	`cos`
\(\vdots\)	\(\vdots\)

After the function is evaluated, Autograd has a graph specifying all operations that were performed on the inputs with respect to which we want to differentiate. This is the computational graph of the function evaluation. To compute the derivative, we simply apply the rules of differentiation to each node in the graph.

\[x\quad\mapsto\quad y=x^2\quad\mapsto\quad z=\sin(y)\]

This is why we need to import Autograd’s alias for the NumPy module: it provides a thin interface layer to support the differentiation.

import autograd.numpy as np

Forward mode and reverse mode

To understand forward mode and reverse mode, we need a slightly more involved example. Suppose we have three functions being composed: \[\ell(x)=h(g(f(x)).\] By applying the chain rule twice, we have the gradient \[\frac{\mathrm{d}\ell}{\mathrm{d}x}=\frac{\mathrm{d}h}{\mathrm{d}g}\frac{\mathrm{d}g}{\mathrm{d}f}\frac{\mathrm{d}f}{\mathrm{d}x}.\]

Thanks to commutativity of multiplication, there are two ways we could evaluate this expression: right-to-left as \[\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)\] or left-to-right as \[\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}.\]

Right-to-left follows the order in which \(f\), \(g\), and \(h\) would be evaluated in the function and is therefore referred to as forward mode. Left-to-right goes against this and is referred to as reverse mode.

Note While these definitions are sufficient for the scalar case, there is more to the story for functions of multiple variables (e.g., vector-valued functions), as we will see in session 2. This will explain why reverse mode (a.k.a. back-propagation) is more commonly used in machine learning methods such as neural networks.

Question

If you want to evaluate both the function and its derivative at the same time, which approach would you choose?

Solution

The chain rule includes values of the component functions as well as their derivatives so forward mode allows for efficient execution.

Example: ODE-constrained optimisation

Consider the scalar ordinary differential equation (ODE) \[ \frac{\mathrm{d}u}{\mathrm{d}t}=f(u),\quad u(0)=u_0, \] where \(t\in[0,T]\) is the time variable, \(T>0\) is the end time, and \(u_0\in\mathbb{R}\) is the initial condition. Given some \(f:\mathbb{R}\rightarrow\mathbb{R}\), we seek to solve the ODE for \(u:[0,T]\rightarrow\mathbb{R}\).

For simplicity, let’s consider the ODE with \(f(u)=u\): \[ \frac{\mathrm{d}u}{\mathrm{d}t}=u,\quad u(0)=1, \]

Optional exercise

Convince yourself that the analytical solution of the ODE is \(u(t)=\mathrm{e}^t\).

Solution

Plugging \(u(t)=\mathrm{e}^t\) into the LHS gives \(\frac{\mathrm{d}u}{\mathrm{d}t}=\mathrm{e}^t=u\), which satisfies the ODE. Checking the initial condition, we have \(u(0)=\mathrm{e}^0=1\), which also satisfies.

The initial condition can be implemented as

def initial_condition():
    """
    Apply the initial condition for the ODE initial value problem
        du/dt = u, u(0) = 1

    :return: initial condition value
    """
    u0 = 1.0
    return u0

ODE example: explicit vs implicit Euler

To solve the ODE we can discretise the \([0,T]\) time domain into timesteps of \(\Delta t > 0\) and find the value of \(u\) at each timestep, \(k\in\mathbb{N}\).

You may be aware of the simplest example of an explicit timestepping method to approximate the solution of the ODE. This is the explicit Euler: \[ \frac{u_{k}-u_{k-1}}{\Delta t}=f(u_{k-1}), \] for \(k\in\mathbb{N}\) and some timestep \(\Delta t>0\).

You’re possibly also aware that the simplest example of an implicit timestepping method to approximate the solution of the ODE is implicit Euler: \[ \frac{u_{k}-u_{k-1}}{\Delta t}=f(u_k). \] In this case, because the left and right hand side both depend on \(u_k\) the equation needs to be rearranged / solved to find \(u_k\) each iteration.

These are special cases of a more general theta-method, \[ \frac{u_{k}-u_{k-1}}{\Delta t}=(1-\theta)f(u_{k-1})+\theta f(u_k), \] where \(\theta\in[0,1]\).

Note The ODE is a model, which involves the derivative of the model solution on the left hand side. In explicit and implicit Euler, this derivative is being approximated using a numerical method, not computed using AD.

ODE example: applying explicit and implicit Euler

Using the method above, our problem reads \[ \frac{u_{k}-u_{k-1}}{\Delta t}=(1-\theta)u_{k-1}+\theta u_k, \] which can be rearranged to give \[ u_{k}=\frac{1+\Delta t(1-\theta)}{1-\Delta t\theta}u_{k-1}. \] We can implement this using a Python function as

def theta_step(u_, dt, theta):
    """
    Take a single iteration of a theta method for solving the ODE initial value problem
        du/dt = u, u(0) = 1

    :arg u_: numerical solution at previous timestep
    :arg dt: timestep length
    :arg theta: parameter to use in timestepping method
    :return: numerical solution at current timestep
    """
    u = u_ * (1.0 + dt * (1.0 - theta)) / (1.0 - dt * theta)
    return u

We can define the \(\theta\)-method in terms of the initial_condition and theta_step functions as

def theta_method(theta):
    """
    Solve the ODE initial value problem
        du/dt = u, u(0) = 1
    using a theta timestepping method, returning the solution trajectory.

    :arg theta: parameter to use in timestepping method
    :return: numerical solution trajectory
    """
    t = 0.0
    dt = 0.1
    end_time = 1.0
    u0 = initial_condition()

    # Timestepping loop
    trajectory = [u0]
    u_ = u0
    while t < end_time - 1.0e-05:
        u = theta_step(u_, dt, theta)
        trajectory.append(u)
        # Update variables for next loop
        u_ = u
        t += dt
    return trajectory

Explicit Euler corresponds to \(\theta=0\) and Implicit Euler corresponds to \(\theta=1\) so we can apply them to the problem with

explicit = theta_method(0.0)
implicit = theta_method(1.0)

times = np.linspace(0, 1, len(explicit))

fig, axes = plt.subplots()
axes.plot(times, np.exp(times), "-", color="k", label="Analytical solution")
axes.plot(times, explicit, "--x", label="Explicit Euler")
axes.plot(times, implicit, ":o", label="Implicit Euler")
axes.legend()
axes.grid()

ODE example: optimisation with gradient descent

As we see from the plot above, the Explicit Euler method tends to underestimate the solution, whereas the Implicit Euler method tends to overestimate it. Let’s try to optimise the value of \(\theta\) to best match the solution using a gradient-based optimisation method. To do that, we first need the gradient. AD enables us to do this automatically.

The optimisation problem we seek to solve is to minimise some error measure \(J\) for the approximation of \(u\) by varying \(\theta\). That is, \[ \min_{\theta\in[0,1]}J(u;\theta). \] where the notation \(J(u;\theta)\) refers to the implicit dependence of the solution approximation \(u\) on \(\theta\).

Note Explicit and implicit Euler are first-order accurate methods. By optimising the \(\theta\) parameter, we can arrive at a second-order accurate method.

Since we know the analytical solution for this problem, we may make an ‘artificial’ choice of objective function such as \[ J(u;\theta)=(u(1)-\mathrm{e}^1)^2, \] where here \(\mathrm{e}^1=2.7182818...\) is the analytical solution at the end time \(t=1\). We can implement this as the Python function

def objective_function(theta):
    """
    Simple cost function evaluating the l2 error at the end time against the analytical solution u(t)=exp(t).

    :arg theta: parameter to use in timestepping method
    :return: l2 error at the end time against the analytical solution
    """
    u = theta_method(theta)[-1]
    e = np.exp(1.0)
    return (u - e) ** 2

Let’s solve this ODE problem with one of the simplest gradient-based optimisation approaches: gradient descent. This amounts to an initial guess \(\theta_0\), followed by iterative updates \[ \theta_{k+1}=\theta_k+\alpha\:p_k, \] where \(\alpha>0\) is the step length and \(p_k\) is the descent direction. For gradient descent, we simply take \[ p_k=-\frac{\mathrm{d}J_k}{\mathrm{d}\theta_k}. \]

A simple implementation of the gradient descent method can be found in the following code block, although there is a missing piece.

Exercise

Crucially, we need to differentiate the objective function with respect to \(\theta\). Rather than trying to do this manually we can use differentiable programming with autograd. Replace the # TODO comment with your implementation.

Solution

dJdtheta = grad(objective_function)(theta)

def gradient_descent(maxiter=1000, gtol=1.0e-05, dtol=1.1, alpha=0.10):
    """
    Function for optimising the theta parameter for a theta timestepping method for solving the ODE
        du/dt = u, u(0)=1
    using gradient descent.

    :arg maxiter: maximum number of iterations
    :arg gtol: fractional tolerance for objective function convergence
    :arg dtol: fractional tolerance for objective function divergence
    :arg alpha: step length
    :return: trajectories of theta and objective function values
    """
    # Start from explicit Euler
    theta = 0.0

    # Create lists for tracking convergence progress
    theta_values = []
    J_values = []
    
    for i in range(maxiter):
        # Compute the objective function and its gradient
        J = objective_function(theta)
        dJdtheta = # TODO: Compute gradient of J with respect to theta
        
        # Record the value of theta and corresponding objective value
        theta_values.append(theta)
        J_values.append(J)
        
        # Convergence and divergence checks
        if i == 0:
            J_init = J
            dJdtheta_init = dJdtheta
        elif abs(dJdtheta / dJdtheta_init) < gtol:
            print(f"Converged in {i+1} iterations due to gradient convergence")
            return theta_values, J_values
        elif abs(J / J_init) > dtol:
            raise RuntimeError(f"Detected divergence after {i+1} iterations")
    
        # Take a step in the descent direction
        p = -dJdtheta
        theta += alpha * p
    raise RuntimeError("Reached maximum iteratons without convergence")

We can then run the gradient descent algorithm and plot it’s progress as follows.

theta_values, J_values = gradient_descent()
theta_opt = theta_values[-1]

fig, axes = plt.subplots(ncols=2, figsize=(12, 5))
axes[0].loglog(J_values, "--", label="Objective function value")
axes[0].legend()
axes[0].grid()
axes[1].plot(theta_values, "--", label=r"$\theta$ value")
axes[1].legend()
axes[1].grid()

This looks promising! Let’s examine the solution trajectory for the optimised \(\theta\) parameter to check it does a better job than Explicit Euler and Implicit Euler.

optimised = theta_method(theta_opt)

fig, axes = plt.subplots()
axes.plot(times, np.exp(times), "-", color="k", label="Analytical solution")
axes.plot(times, explicit, "--x", label="Explicit Euler")
axes.plot(times, implicit, ":o", label="Implicit Euler")
axes.plot(times, optimised, "-.^", label=rf"Optimised ($\theta={theta_opt:.4f}$)")
axes.legend()
axes.grid()

As we might hope, the optimised value of \(\theta\) gives a much better approximation.

Computing higher-order derivatives

Now let’s try computing higher order derivatives using the hyperbolic tangent function,

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

This can be implemented as the Python function:

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

Exercise

Below is the same code for plotting this function over the range \([-7,7]\). Add code for printing its first four derivatives over the same range.

Hint: To compute gradients of a scalar-valued function over a range of values, make use of elementwise_grad rather than grad. Call help(grad) to see its docstring.

Solution

from autograd import elementwise_grad as egrad

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

fig, axes = plt.subplots()
axes.plot(x, tanh(x), label=r"$\tanh(x)$")
axes.plot(x, egrad(tanh))(x), label="First derivative")
axes.plot(x, egrad(egrad(tanh))(x), label="Second derivative")
axes.plot(x, egrad(egrad(egrad(tanh)))(x), label="Third derivative")
axes.plot(x, egrad(egrad(egrad(egrad(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 autograd import elementwise_grad as egrad

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();

Neural network example

While we haven’t covered the case of functions with more than one input or output, it is still illustrative to see the power of the approach in practice for the kind of problem you might want to solve in your work.

A neural network is a piece of machine learning technology that can be used to make predictions based on data. As suggested by the name, they are designed to mimick the activation of neurons in the human brain. As a result, like humans, neural networks significantly out-perform traditional computational approaches for tasks such as image recognition. In this worked example, we will demonstrate how differentiable programming is an enabling technology in achieving that.

The core idea is to consider a graph such as the one shown below and to weight the nodes (“neurons”) and the connections between them and interpret traversing the network as a series of matrix-vector multiplications, where the matrix and vector values are determined by these weights. Input data is provided to the input nodes on the left-hand-side - in this case corresponding to an image representing a numerical digit - and these data are propagated through the network to predict which digit it corresponds to.

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

To achieve a reasonable result, the neural network needs to be trained. This amounts to tuning the weights such that the output more often matches expectations. For example, an image of the number 3 should be correctly identified as such. This is achieved in much the same way as for the ODE-constrained optimisation example: we define an objective (a.k.a. loss) function and apply a gradient-based optimisation method, with the gradient computed using AD.

First, download the training and testing data set, known as MNIST.

import autograd.numpy as np
import array
import gzip
import os
import struct
from urllib.request import urlretrieve


def download(url, filename):
    """
    Create a data directory and download a file into it from a URL.

    :arg url: URL to download the file from
    :arg filename: the name of the file
    """
    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 mnist():
    """
    Download the MNIST training data as .gz files.

    :return: arrays of images and labels for the training and testing sets.
    """
    base_url = "https://storage.googleapis.com/cvdf-datasets/mnist/"

    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)

    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")

    return train_images, train_labels, test_images, test_labels


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

    :return: number of images and arrays of normalised images and labels
    """
    partial_flatten = lambda x: np.reshape(x, (x.shape[0], np.prod(x.shape[1:])))
    one_hot = lambda x, k: np.array(x[:, None] == np.arange(k)[None, :], dtype=int)
    train_images, train_labels, test_images, test_labels = mnist()
    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

print("Loading training data...")
N, train_images, train_labels, test_images, test_labels = load_mnist()
print(f"Training dataset size: {N}")
print(f"Testing dataset size: {len(test_images)}")

Next, define the neural network architecture and the components of the objective function.

import autograd.numpy.random as npr
from autograd.misc.flatten import flatten
from autograd.misc.optimizers import adam
from autograd.scipy.special import logsumexp


def init_random_params(scale, layer_sizes, rs=npr.RandomState(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
    """
    return [
        (
            scale * rs.randn(m, n),  # weight matrix
            scale * rs.randn(n),     # bias vector
        )  
        for m, n in zip(layer_sizes[:-1], layer_sizes[1:])
    ]


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:
        outputs = np.dot(inputs, W) + b
        inputs = np.tanh(outputs)
    return np.exp(outputs - logsumexp(outputs, axis=1, keepdims=True)) # softmax


def l2_norm(params):
    """
    Computes l2 norm of params by flattening them into a vector.

    :arg params: neural network parameters
    :return: l2 norm of the parameters
    """
    flattened, _ = flatten(params)
    return np.dot(flattened, flattened)


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
    """
    log_prior = -L2_reg * l2_norm(params)
    probs = neural_net_predict(params, inputs)
    log_lik = np.sum(np.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 = np.argmax(targets, axis=1)
    predicted_class = np.argmax(neural_net_predict(params, inputs), axis=1)
    return np.mean(predicted_class == target_class)


# Model parameters
param_scale = 0.1
layer_sizes = [784, 200, 100, 10]
L2_reg = 1.0

init_params = init_random_params(param_scale, layer_sizes)

We can inspect the number of parameters in each layer as follows:

for i, param in enumerate(init_params):
    W, b = param
    print(f"layer: {i}, weights: {W.shape}, biases: {b.shape}")

To verify that the network gives garbage before being trained, let’s take a look at some of the test data images and check whether the digits they represent are correctly identified.

def plot_images(
    images,
    ax,
    ims_per_row=5,
    padding=5,
    digit_dimensions=(28, 28),
    cmap=matplotlib.cm.binary,
    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
    cax = ax.matshow(concat_images, cmap=cmap, vmin=vmin, vmax=vmax)
    plt.xticks(np.array([]))
    plt.yticks(np.array([]))
    return cax

fig, axes = plt.subplots()
plot_images(test_images[100:105], axes);

neural_net_predict(init_params, test_images[100:105]).argmax(axis=1)

Finally, solve the optimisation problem. In this example, we use the adam optimisation method rather than the simple gradient descent approach because it is better suited to ‘noisy’ machine learning problems.

Exercise

You will also need to differentiate the objective function with respect to the neural network parameters. Replace the # TODO comment with your implementation of the gradient function.

Solution

objective_grad = grad(objective)

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

num_batches = int(np.ceil(len(train_images) / batch_size))

def batch_indices(iteration):
    idx = iteration % num_batches
    return slice(idx * batch_size, (idx + 1) * batch_size)

# Define training objective
def objective(params, iteration):
    idx = batch_indices(iteration)
    return -log_posterior(params, train_images[idx], train_labels[idx], L2_reg)

# Get gradient of objective using autograd.
objective_grad = # TODO

def print_perf(params, iteration, gradient):
    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}")

# The optimizers provided can optimize lists, tuples, or dicts of parameters.
optimized_params = adam(
    objective_grad,
    init_params,
    step_size=step_size,
    num_iters=num_epochs * num_batches,
    callback=print_perf,
)

Whilst it would be possible to do the optimisation without differentiable programming, the task of deriving derivatives by hand requires manual effort and is a tedious and error-prone task.

To verify that the network was trained as required, let’s take a look at some of the test data images and check that the digits they represent are correctly identified.

fig, axes = plt.subplots()
plot_images(test_images[100:105], axes);

neural_net_predict(optimized_params, test_images[100:105]).argmax(axis=1)

Success! Hopefully you are now convinced that AD is of crucial importance for machine learning and other scientific software approaches.

The example shown above could also be achieved in PyTorch - see the course on Wednesday morning.

Other operator overloading AD tools

There are many other tools that follow a similar approach. An extensive list can be found at https://autodiff.org/?module=Tools but here are a few notable ones for the programming languages most relevant to ICCS. (Bold indicates ICCS-led projects.)

Python

C/C++

ADIC
ADOL-C
Torch Autograd
CoDiPack
Sacado
Enzyme
dco/c++
about two dozen overall!

Fortran

Differentia
FTorch
Enzyme
and lots of abandonware…

Julia

Summary and outlook

In today’s session we:

Got a brief history of automatic differentiation.
Learnt about the operator overloading approach.
Learnt about forward mode and reverse mode applied to scalar-valued functions.
Tried out the Autograd operator overloading AD tool applied to some test problems.
Verified the derivatives produced by Autograd using the Taylor test.

In tomorrow’s session (11:00-12:30) 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.

European workshop on Automatic differentiation

Figure 4: 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.

References

Autograd Tutorial
R. E. Wengert. A simple automatic derivative evaluation program (1964). Communications of the ACM, 7(8):463–464, doi.org:10.1145/355586.364791.
S. Linnainmaa. Taylor expansion of the accumulated rounding error. BIT, 16(2):146–160, 1976, doi:10.1007/BF01931367.
B. Speelpenning. Compiling fast partial derivatives of functions given by algorithms. University of Illinois, 1980, doi:10.2172/5254402.
A. Griewank. Achieving logarithmic growth of temporal and spatial complexity in reverse automatic differentiation. Optimization Methods & Software, 1:35–54, 1992, doi:10.1080/10556789208805505.