1.5  Automatic Differentiation

Recall from Section 1.4 that derivatives drive all the optimization algorithms that we will use to train deep networks. Even when the individual rules are elementary, applying them by hand becomes tedious and error-prone as models grow more complex.

Modern deep learning frameworks perform this calculation using automatic differentiation (often shortened to autograd). As we pass data through each successive function, the framework builds a computational graph that tracks how each value depends on others. To calculate derivatives, automatic differentiation works backwards through this graph applying the chain rule. The computational algorithm for applying the chain rule in this fashion is called backpropagation.

Autograd has a long history: the earliest references date back over half a century (Wengert 1964), and reverse mode, the variant that powers modern backpropagation, was developed by Linnainmaa (1970) . Section 25.3 recounts this history in full. We begin with the basic automatic-differentiation interface and then examine how computational graphs are controlled.

import torch
import tensorflow as tf
from jax import numpy as jnp
from mxnet import autograd, np, npx
npx.set_np()

1.5.1 Mechanics

We begin with the basic workflow (attach a gradient, record a computation, run the backward pass), first on a scalar-valued function, then on vector-valued ones.

1.5.1.1 A Simple Function

Consider differentiating the function \(y = 2\mathbf{x}^{\top}\mathbf{x}\) with respect to the column vector \(\mathbf{x}\). To start, we assign x an initial value.

x = torch.arange(4.0)
x
tensor([0., 1., 2., 3.])
x = tf.range(4, dtype=tf.float32)
x
<tf.Tensor: shape=(4,), dtype=float32, numpy=array([0., 1., 2., 3.], dtype=float32)>
x = jnp.arange(4.0)
x
Array([0., 1., 2., 3.], dtype=float32)
x = np.arange(4.0)
x
array([0., 1., 2., 3.])

Before we calculate the gradient of \(y\) with respect to \(\mathbf{x}\), we need a place to store it. In general, we avoid allocating new memory every time we take a derivative because deep learning requires successively computing derivatives with respect to the same parameters a great many times, and repeated allocation would add unnecessary overhead. The gradient of a scalar-valued function with respect to a vector \(\mathbf{x}\) is vector-valued with the same shape as \(\mathbf{x}\).

Before we calculate the gradient of \(y\) with respect to \(\mathbf{x}\), we need a place to store it. In general, we avoid allocating new memory every time we take a derivative because deep learning requires successively computing derivatives with respect to the same parameters a great many times, and repeated allocation would add unnecessary overhead. The gradient of a scalar-valued function with respect to a vector \(\mathbf{x}\) is vector-valued with the same shape as \(\mathbf{x}\).

Before we calculate the gradient of \(y\) with respect to \(\mathbf{x}\), we need a place to store it. In general, we avoid allocating new memory every time we take a derivative because deep learning requires successively computing derivatives with respect to the same parameters a great many times, and repeated allocation would add unnecessary overhead. The gradient of a scalar-valued function with respect to a vector \(\mathbf{x}\) is vector-valued with the same shape as \(\mathbf{x}\).

# Can also create x = torch.arange(4.0, requires_grad=True)
x.requires_grad_(True)
x.grad  # The gradient is None by default
x = tf.Variable(x)
# We allocate memory for a tensor's gradient by invoking `attach_grad`
x.attach_grad()
# After we calculate a gradient taken with respect to `x`, we will be able to
# access it via the `grad` attribute, whose values are initialized with 0s
x.grad
array([0., 0., 0., 0.])

We now calculate our function of x and assign the result to y.

y = 2 * torch.dot(x, x)
y
tensor(28., grad_fn=<MulBackward0>)
# Record all computations onto a tape
with tf.GradientTape() as t:
    y = 2 * tf.tensordot(x, x, axes=1)
y
<tf.Tensor: shape=(), dtype=float32, numpy=28.0>
y = lambda x: 2 * jnp.dot(x, x)
y(x)
Array(28., dtype=float32)
# Our code is inside an `autograd.record` scope to build the computational
# graph
with autograd.record():
    y = 2 * np.dot(x, x)
y
array(28.)

Recording the operations gives the framework a computational graph, shown in Figure 1.5.1. Its nodes are operations and its edges carry intermediate values.

Figure 1.5.1: The computational graph for \(y = 2\mathbf{x}^\top\mathbf{x}\). The forward pass (black) flows from \(\mathbf{x}\) to \(y\); reverse-mode automatic differentiation walks the same graph backward (blue), multiplying the local derivative at each node via the chain rule to accumulate \(\partial y / \partial \mathbf{x} = 4\mathbf{x}\).

The forward pass evaluates the graph from \(\mathbf{x}\) to \(y\); to obtain the gradient, automatic differentiation then traverses it in reverse, multiplying the local derivatives along the way. We unpack computational graphs and backpropagation in full in Section 4.3, and the underlying mathematics (both modes of automatic differentiation and their costs) is developed in Section 25.3; for now we use the resulting gradients.

We can now take the gradient of y with respect to x by calling its backward method. Next, we can access the gradient via x’s grad attribute.

We can now calculate the gradient of y with respect to x by calling the gradient method.

We can now take the gradient of y with respect to x by passing through the grad transform.

We can now take the gradient of y with respect to x by calling its backward method. Next, we can access the gradient via x’s grad attribute.

y.backward()
x.grad
tensor([ 0.,  4.,  8., 12.])
x_grad = t.gradient(y, x)
x_grad
<tf.Tensor: shape=(4,), dtype=float32, numpy=array([ 0.,  4.,  8., 12.], dtype=float32)>
from jax import grad
# The `grad` transform returns a Python function that
# computes the gradient of the original function
x_grad = grad(y)(x)
x_grad
Array([ 0.,  4.,  8., 12.], dtype=float32)
y.backward()
x.grad
array([ 0.,  4.,  8., 12.])

We already know that the gradient of the function \(y = 2\mathbf{x}^{\top}\mathbf{x}\) with respect to \(\mathbf{x}\) should be \(4\mathbf{x}\). We can now verify that the automatic gradient computation and the expected result are identical.

x.grad == 4 * x
tensor([True, True, True, True])
x_grad == 4 * x
<tf.Tensor: shape=(4,), dtype=bool, numpy=array([ True,  True,  True,  True])>
x_grad == 4 * x
Array([ True,  True,  True,  True], dtype=bool)
x.grad == 4 * x
array([ True,  True,  True,  True])

For the new objective \(y=\sum_i x_i\), PyTorch adds the new derivative to the existing gradient buffer unless it is cleared first. This behavior supports optimizing the sum of multiple objective functions. To reset the gradient buffer, we can call x.grad.zero_() as follows:

For the new objective \(y=\sum_i x_i\), GradientTape.gradient returns a new gradient value rather than accumulating into a tensor-owned buffer.

For the new objective \(y=\sum_i x_i\), MXNet replaces the recorded gradient with the new value.

x.grad.zero_()  # Reset the gradient
y = x.sum()
y.backward()
x.grad
tensor([1., 1., 1., 1.])
with tf.GradientTape() as t:
    y = tf.reduce_sum(x)
t.gradient(y, x)  # Overwritten by the newly calculated gradient
<tf.Tensor: shape=(4,), dtype=float32, numpy=array([1., 1., 1., 1.], dtype=float32)>
y = lambda x: x.sum()
grad(y)(x)
Array([1., 1., 1., 1.], dtype=float32)
with autograd.record():
    y = x.sum()
y.backward()
x.grad  # Overwritten by the newly calculated gradient
array([1., 1., 1., 1.])

1.5.1.2 Backward for Non-Scalar Variables

When y is a vector, the most natural representation of the derivative of y with respect to a vector x is a matrix called the Jacobian that contains the partial derivatives of each component of y with respect to each component of x. Likewise, for higher-order y and x, the result of differentiation could be an even higher-order tensor.

While Jacobians do show up in some advanced machine learning techniques, more commonly we sum the gradients of each component of y with respect to the full vector x, yielding a vector of the same shape as x. For example, we often have a vector representing the value of our loss function calculated separately for each example among a batch of training examples. Here, we sum the gradients computed individually for each example.

Because deep learning frameworks vary in how they interpret gradients of non-scalar tensors, PyTorch takes some steps to avoid confusion. Invoking backward on a non-scalar elicits an error unless we tell PyTorch how to reduce the object to a scalar. More formally, we need to provide some vector \(\mathbf{v}\) such that backward will compute \(\mathbf{v}^\top \partial_{\mathbf{x}} \mathbf{y}\) rather than \(\partial_{\mathbf{x}} \mathbf{y}\). This argument is named gradient because the vector \(\mathbf{v}\) is the gradient arriving from the rest of a larger computation, as will become clear when we study backpropagation in Section 4.3. For a more detailed description, see the PyTorch documentation on the gradient argument to Tensor.backward.

By default, TensorFlow returns the gradient of the sum. In other words, rather than returning the Jacobian \(\partial_{\mathbf{x}} \mathbf{y}\), it returns the gradient of the sum \(\partial_{\mathbf{x}} \sum_i y_i\).

MXNet handles this problem by reducing all tensors to scalars by summing before computing a gradient. In other words, rather than returning the Jacobian \(\partial_{\mathbf{x}} \mathbf{y}\), it returns the gradient of the sum \(\partial_{\mathbf{x}} \sum_i y_i\).

x.grad.zero_()
y = x * x
y.backward(gradient=torch.ones(len(y)))  # Equivalently: y.sum().backward()
x.grad
tensor([0., 2., 4., 6.])
with tf.GradientTape() as t:
    y = x * x
t.gradient(y, x)  # Same as y = tf.reduce_sum(x * x)
<tf.Tensor: shape=(4,), dtype=float32, numpy=array([0., 2., 4., 6.], dtype=float32)>
y = lambda x: x * x
# grad is only defined for scalar output functions
grad(lambda x: y(x).sum())(x)
Array([0., 2., 4., 6.], dtype=float32)
with autograd.record():
    y = x * x  
y.backward()
x.grad  # Equals the gradient of y = sum(x * x)
array([0., 2., 4., 6.])

1.5.2 Controlling the Graph

Sometimes the graph the framework records is not the graph we want to differentiate. The next two subsections show how to prune it by detaching individual intermediate results, and how to switch recording off altogether.

1.5.2.1 Detaching Computation

Sometimes, we wish to move some calculations outside of the recorded computational graph. For example, say that we use the input to create some auxiliary intermediate terms for which we do not want to compute a gradient. In this case, we need to detach the respective computational graph from the final result. For example, suppose that z = x * y and y = x * x but we want to focus on the direct influence of x on z rather than the influence conveyed via y. In this case, we can create a new variable u that takes the same value as y but whose computational history is not connected to the new graph. Thus u has no ancestors in the graph and gradients do not flow through u to x. Now consider z = x * u. Because u is treated as a constant equal to \(x^2\), the gradient is \(\partial z / \partial x = u = x^2\). Had we not detached, so that z = x * (x * x) \(= x^3\), we would instead have obtained \(\partial z / \partial x = 3x^2\).

x.grad.zero_()
y = x * x
u = y.detach()
z = u * x

z.sum().backward()
x.grad == u
tensor([True, True, True, True])
# Set persistent=True to preserve the compute graph. 
# This lets us run t.gradient more than once
with tf.GradientTape(persistent=True) as t:
    y = x * x
    u = tf.stop_gradient(y)
    z = u * x

x_grad = t.gradient(z, x)
x_grad == u
<tf.Tensor: shape=(4,), dtype=bool, numpy=array([ True,  True,  True,  True])>
import jax

y = lambda x: x * x
# jax.lax primitives are Python wrappers around XLA operations
u = jax.lax.stop_gradient(y(x))
z = lambda x: u * x

grad(lambda x: z(x).sum())(x) == u
Array([ True,  True,  True,  True], dtype=bool)
with autograd.record():
    y = x * x
    u = y.detach()
    z = u * x
z.backward()
x.grad == u
array([ True,  True,  True,  True])

Although this procedure detaches y’s ancestors from the graph leading to z, the computational graph leading to y persists and thus we can calculate the gradient of y with respect to x.

x.grad.zero_()
y.sum().backward()
x.grad == 2 * x
tensor([True, True, True, True])
t.gradient(y, x) == 2 * x
<tf.Tensor: shape=(4,), dtype=bool, numpy=array([ True,  True,  True,  True])>
grad(lambda x: y(x).sum())(x) == 2 * x
Array([ True,  True,  True,  True], dtype=bool)
y.backward()
x.grad == 2 * x
array([ True,  True,  True,  True])

1.5.2.2 Turning Off Gradient Tracking

Recording operations for a backward pass costs time and memory. When we only need a value (at prediction time, or while updating parameters by hand), we can skip the bookkeeping entirely.

Wrap the computation in a torch.no_grad() block (or decorate a function with @torch.no_grad()). The result still shares data with x, but it is not attached to the graph, so no gradient can flow through it.

TensorFlow only records operations executed inside a tf.GradientTape (a tape is the recorded list of executed operations), so any computation outside a tape is already untracked. To pause recording within a tape, use tape.stop_recording().

JAX never records gradients implicitly: differentiation occurs only after a transform such as grad is applied. Thus, users opt in to differentiation.

MXNet only builds a graph inside an autograd.record() block, so ordinary computation already carries no gradient bookkeeping. To suspend tracking within a recording scope, wrap the code in autograd.pause().

with torch.no_grad():
    y = 2 * torch.dot(x, x)
y.requires_grad  # False: y is detached from the graph
False
# Outside any GradientTape, nothing is recorded
y = 2 * tf.tensordot(x, x, axes=1)
y
<tf.Tensor: shape=(), dtype=float32, numpy=28.0>
# No graph is built unless we ask for it via a transform like `grad`
y = 2 * jnp.dot(x, x)
y
Array(28., dtype=float32)
with autograd.record():
    with autograd.pause():
        y = 2 * np.dot(x, x)  # not recorded: no gradient will flow through y
y
array(28.)

This untracked mode is the default for inference and evaluation throughout the rest of the book.

1.5.3 Beyond the Basics

Automatic differentiation extends beyond the fixed formulas considered so far: it handles arbitrary control flow, derivatives of derivatives, and even lets us choose the direction in which the graph is traversed.

1.5.3.1 Gradients and Python Control Flow

The preceding examples used a fixed sequence of operations such as z = x * x * x. Programs can instead use auxiliary variables and choose branches from intermediate results. One benefit of using automatic differentiation is that even if building the computational graph of a function required Python control flow (e.g., conditionals, loops, and arbitrary function calls), we can still calculate the gradient of the resulting variable. To illustrate this, consider the following code snippet where the number of iterations of the while loop and the evaluation of the if statement both depend on the value of the input a.

def f(a):
    b = a * 2
    while b.norm() < 1000:
        b = b * 2
    if b.sum() > 0:
        c = b
    else:
        c = 100 * b
    return c
def f(a):
    b = a * 2
    while tf.norm(b) < 1000:
        b = b * 2
    if tf.reduce_sum(b) > 0:
        c = b
    else:
        c = 100 * b
    return c
def f(a):
    b = a * 2
    while jnp.linalg.norm(b) < 1000:
        b = b * 2
    if b.sum() > 0:
        c = b
    else:
        c = 100 * b
    return c
def f(a):
    b = a * 2
    while np.linalg.norm(b) < 1000:
        b = b * 2
    if b.sum() > 0:
        c = b
    else:
        c = 100 * b
    return c

Below, we call this function, passing in a random value, as input. Since the input is a random variable, we do not know what form the computational graph will take. However, whenever we execute f(a) on a specific input, we realize a specific computational graph and can subsequently run backward.

a = torch.randn(size=(), requires_grad=True)
d = f(a)
d.backward()
a = tf.Variable(tf.random.normal(shape=()))
with tf.GradientTape() as t:
    d = f(a)
d_grad = t.gradient(d, a)
d_grad
<tf.Tensor: shape=(), dtype=float32, numpy=4096.0>
from jax import random
a = random.normal(random.key(1), ())
d = f(a)
d_grad = grad(f)(a)
a = np.random.normal()
a.attach_grad()
with autograd.record():
    d = f(a)
d.backward()

Even though our function f is, for demonstration purposes, a bit contrived, its dependence on the input is quite simple: it is a linear function of the scalar a with piecewise defined scale. As such, f(a) / a is a constant and, moreover, it needs to match the gradient of f(a) with respect to a.

a.grad == d / a
tensor(True)
d_grad == d / a
<tf.Tensor: shape=(), dtype=bool, numpy=True>
d_grad == d / a
Array(True, dtype=bool)
a.grad == d / a
array(True)

Dynamic control flow occurs frequently in deep learning. For instance, when processing text, the computational graph depends on the length of the input. Automatic differentiation records or traces the operations for the realized computation rather than requiring a separate derivative for every possible execution path.

1.5.3.2 Higher-Order Derivatives

Occasionally we need the derivative of a derivative: the curvature of a function, or the Hessian–vector products (products of the matrix of second derivatives, the Hessian, with a vector) used by some optimizers. Autograd can differentiate through a gradient computation. Take \(f(x) = x^3\), for which \(f'(x) = 3x^2\) and \(f''(x) = 6x\).

Pass create_graph=True so the first gradient is itself a differentiable function of x, then differentiate again.

Nest two GradientTapes: the outer tape differentiates the gradient computed under the inner tape.

grad returns a function, so we apply it twice.

Higher-order gradients in MXNet require explicitly retaining the graph of the first derivative. The mathematics is framework-agnostic and developed in Section 25.3.

x3 = torch.tensor(2.0, requires_grad=True)
dy = torch.autograd.grad(x3 ** 3, x3, create_graph=True)[0]  # 3x^2 = 12
d2y = torch.autograd.grad(dy, x3)[0]                          # 6x  = 12
dy, d2y
(tensor(12., grad_fn=<MulBackward0>), tensor(12.))
x3 = tf.Variable(2.0)
with tf.GradientTape() as outer:
    with tf.GradientTape() as inner:
        y = x3 ** 3
    dy = inner.gradient(y, x3)   # 3x^2 = 12
d2y = outer.gradient(dy, x3)     # 6x  = 12
dy, d2y
(<tf.Tensor: shape=(), dtype=float32, numpy=12.0>,
 <tf.Tensor: shape=(), dtype=float32, numpy=12.0>)
f = lambda x: x ** 3
dy = grad(f)(2.0)            # 3x^2 = 12
d2y = grad(grad(f))(2.0)    # 6x  = 12
dy, d2y
(Array(12., dtype=float32, weak_type=True),
 Array(12., dtype=float32, weak_type=True))

1.5.3.3 Forward versus Reverse Mode

Automatic differentiation can traverse the computational graph in either direction. Reverse mode, the variant we have used so far and the engine behind backpropagation, sweeps from the output back to the inputs, yielding the gradient of a single scalar with respect to all inputs in one pass. Forward mode sweeps the other way, propagating derivatives from one input outward to every output.

The choice is about cost, and a counting argument settles it. For a function with \(n\) inputs and \(m\) outputs, filling the full matrix of derivatives takes one reverse sweep per output (\(m\) sweeps) or one forward sweep per input (\(n\) sweeps), each sweep costing about as much as one evaluation of the function. A training loss is a single scalar (\(m = 1\)) depending on millions of parameters (\(n\) huge), so reverse mode delivers the entire gradient for the price of roughly one extra forward pass. Forward mode wins in the opposite regime (few inputs, many outputs), is useful for per-input sensitivities, and can participate in efficient Hessian–vector-product constructions, as in the Julia package ForwardDiff.jl (Revels et al. 2016). The exercises explore this trade-off further, and Section 25.3 derives both modes and their costs in full.

1.5.4 Discussion

Automatic differentiation frees practitioners from deriving gradients by hand, and it makes it practical to train models for which pen and paper gradient computations would be prohibitively time consuming. While we use autograd to optimize models (in a statistical sense), the optimization of autograd libraries themselves (in a computational sense) is a rich subject that matters to framework designers. Here, tools from compilers and graph manipulation are used to compute results quickly and with modest memory.

The basic workflow is: (i) attach gradients to those variables with respect to which we desire derivatives; (ii) record the computation of the target value; (iii) execute the backpropagation function; and (iv) access the resulting gradient.

1.5.5 Exercises

  1. Cost of second derivatives. ● The counting argument in Section 1.5.3.3 prices the gradient of a scalar loss with \(n\) parameters at roughly one extra forward pass. Use the same argument to estimate, in units of function evaluations, the cost of the full Hessian (the \(n \times n\) matrix of second derivatives) and the cost of a single Hessian–vector product. Explain the difference between the two costs.
  1. [code] Repeated backward passes. For an x with requires_grad=True, compute y = 2 * torch.dot(x, x), call y.backward(), and then call y.backward() a second time. What happens? Reset x.grad, recompute y, and repeat both calls, this time passing retain_graph=True to the first one. Explain what the backward pass discards by default and why, and what the second call leaves in x.grad.
  1. [code] Repeated backward passes. Record y = 2 * tf.tensordot(x, x, axes=1) under with tf.GradientTape() as t:, call t.gradient(y, x), and then call it a second time. What happens? A tape constructed with persistent=True keeps its recorded operations so that t.gradient can be called repeatedly. Explain what this costs in memory during training, and when the memory held by a persistent tape is released.
  1. [code] Repeated backward passes. Call the function returned by jax.grad twice on the same input. Why does the second call behave no differently from the first? What does this tell you about where JAX keeps the information needed for differentiation?
  1. [code] Repeated backward passes. For an x with an attached gradient, compute y = 2 * np.dot(x, x) inside autograd.record(), call y.backward(), and then call y.backward() a second time. What happens? Record y again and repeat both calls, this time passing retain_graph=True to the first one. Explain what the backward pass discards by default and why, and what the second call leaves in x.grad.
  1. [code] From scalar to vector input. Replace the scalar input a of the function f defined in Section 1.5.3.1 by a random vector of length 5, and differentiate the sum of the entries of d = f(a) with respect to a.
    1. Express f(a) in terms of a, and give the Jacobian of f at a and the gradient of the sum of the entries of d.
    2. Compute this gradient with automatic differentiation and check whether it still equals d / a elementwise.
    3. Identify the inputs a at which f jumps, and state what automatic differentiation returns at an input next to such a point.
  2. [code] Manual tracing. Let \(f(x, y, z) = (x + y) \cdot z\) with \(x = -2\), \(y = 5\), and \(z = -4\).
    1. Draw the computational graph of \(f\) with one node per operation, compute the forward pass by hand, and trace \(\partial f/\partial x\), \(\partial f/\partial y\), and \(\partial f/\partial z\) backward from \(f\) through the multiplication and addition nodes (reverse mode).
    2. Compute the same three derivatives in forward mode, propagating the derivative with respect to one input at a time from the inputs to \(f\). Compare the number of sweeps each mode needs and the intermediate values each must keep in memory.
    3. Verify every value using automatic differentiation.
    Adapted from Stanford CS231n, “An Exercise in Backpropagation”.
  1. [code] Detached intermediates. Let y = x * x and u = y.detach(), and consider z = u * u, z = y * u, and z = y * y. For each z, recompute y and u, reset x.grad, and call backward on the sum of the components of z. Before running any code, predict the outcome of each call, including the resulting x.grad. Check your predictions, and explain the ratio between the gradients for z = y * u and z = y * y.
  2. [code] Manual parameter updates. Let w be a tensor with requires_grad=True, compute loss = (w * w).sum(), and call loss.backward(). Perform the update w = w - 0.1 * w.grad and inspect w.requires_grad, w.grad_fn, w.is_leaf, and w.grad. Explain what the update has done to the graph, predict what a second forward and backward pass leaves in w.grad, and check. Then perform the update in place inside a torch.no_grad() block and compare.
  1. [code] Detached intermediates. Let y = x * x and u = tf.stop_gradient(y) inside a persistent tf.GradientTape t. Before running any code, predict t.gradient(z, x) for each of z = u * u, z = y * u, and z = y * y, all computed inside t. Check your predictions, and explain the ratio between the results for z = y * u and z = y * y.
  2. [code] Manual parameter updates. Let w be a tf.Variable, compute loss = tf.reduce_sum(w * w) under a tf.GradientTape t, and obtain g = t.gradient(loss, w).
    1. Perform the update w = w - 0.1 * g and inspect type(w). Predict what gradient returns for a second forward pass under a new tape, check, and explain the result.
    2. Start again from a tf.Variable, perform the update with w.assign_sub(0.1 * g), and repeat the second forward pass. Explain why a training loop updates variables this way.
  1. [code] Detached intermediates. Let y = lambda x: x * x and u = jax.lax.stop_gradient(y(x)). Before running any code, predict grad(lambda x: z(x).sum())(x) for each of z = lambda x: u * u, z = lambda x: y(x) * u, and z = lambda x: y(x) * y(x). Check your predictions, and explain the ratio between the results for z = lambda x: y(x) * u and z = lambda x: y(x) * y(x).
  2. [code] Manual parameter updates. Let w = jnp.arange(4.0) and g = lambda x: (w * x).sum(), which reads w from the enclosing scope. Predict grad(g)(x), check, and state with respect to which array grad differentiates. Then rewrite g so that grad returns the gradient of (w * x).sum() with respect to w, use it for one update of w with learning rate 0.1, and explain what this implies for how a JAX training loop must pass the parameters to the loss function.
  1. [code] Detached intermediates. Let y = x * x and u = y.detach(), and consider z = u * u, z = y * u, and z = y * y. For each z, compute y, u, and z afresh inside an autograd.record() scope and call z.backward(). Before running any code, predict the outcome of each call, including the resulting x.grad. Check your predictions, and explain the ratio between the gradients for z = y * u and z = y * y.
  2. [code] Manual parameter updates. Let w be an array with an attached gradient, compute loss = (w * w).sum() inside autograd.record(), and call loss.backward(). Perform the update w = w - 0.1 * w.grad outside the recording scope and inspect w.grad. Explain what the rebinding has done to the buffer attached by attach_grad, and predict what a second recorded forward pass followed by backward does. Then perform the update as w[:] = w - 0.1 * w.grad instead and compare.