We can't find the internet
Attempting to reconnect
Something went wrong!
Attempting to reconnect
Compute Jacobian
A Jacobian records how each output of a vector-valued function changes with each input. Rows correspond to outputs; columns correspond to inputs.
Your task
Implement compute_jacobian(x) for a floating-point vector of shape (2,).
The function to differentiate is:
f(x) = [x[0] * x[1], x[0] + x[1]**2]
Return a tensor of shape (2, 2) with x's dtype and device. Leave x unchanged.
Entry J[i, j] is the derivative of output i with respect to input j:
J = [[x[1], x[0]],
[1, 2*x[1]]]
For x=[3, 4], return [[4, 3], [1, 8]]. In the second row, changing x[0]
changes the second output at rate 1, while changing x[1] changes it at rate 8.
Do not transpose this layout.
Define f using tensor operations and stack its two scalar outputs. Then use
torch.autograd.functional.jacobian(f, x) or jax.jacobian(f)(x).
The formula above lets you check the result while practicing automatic
differentiation.
Differentiating the sum of f's outputs gives only one gradient vector: it adds the Jacobian's rows together. This task needs both rows separately, not that summed gradient and not the original two function values.
Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.
Compute derivatives by following the operations
Automatic differentiation computes derivatives by applying the chain rule to a program's operations. It does not estimate them by evaluating nearby inputs, as finite differences do, or rearrange the whole expression using symbolic algebra.
For f(g(x)), the chain rule combines the derivative of g with the derivative
of f at g(x). Repeating this process lets a framework differentiate a long
sequence of operations.
Forward mode follows an input change
For a function with n inputs and m outputs, the Jacobian J is an (m, n)
matrix of derivatives. Entry (i, j) describes how output i changes with input
j.
Forward-mode differentiation computes Jv: the output's rate of change when the
input moves in direction v. For example, if f(x) = sum(x_i^2), then at
x = [1, 2] the gradient is [2, 4]. In direction v = [3, 0], the rate of
change is 2 * 3 + 4 * 0 = 6.
This is a directional derivative. It gives one Jacobian column only when v
selects one input coordinate, such as [1, 0].
Reverse mode works backward from an output
Reverse mode computes J^T u, where u weights the outputs. For a scalar loss,
choosing u = 1 gives the gradient with respect to all inputs. This makes reverse
mode useful for training: a model can have many parameters but just one loss.
Forward mode can be useful when you need sensitivity to only a few input directions. Neither mode needs to construct the full Jacobian for these products. That can save substantial memory. Higher-order differentiation can similarly compute a Hessian-vector product without storing the full matrix of second derivatives.
Changing gradient flow changes what is differentiated
A stop-gradient operation keeps a value in the forward calculation but treats it as constant during differentiation. A custom derivative supplies a derivative rule explicitly. It should match the intended function unless the exercise asks for a surrogate, such as an approximation used during training.
Your code, saved here
Sign in to access your accepted solutions and saved draft history. You can keep reading the lesson without an account.
Sign in to your accountStuck?
JAX reference solution
Sign in to track your attempts toward unlocking the reference solution.