The Jacobian Matrix

The gradient handles functions that return one number. Most layers return a vector, so their derivative is a matrix — the Jacobian. Autograd rarely builds it explicitly, but knowing its shape explains what a vector-Jacobian product is doing.

The idea

For $f: \mathbb{R}^n \to \mathbb{R}^m$, the Jacobian is the $m \times n$ matrix

Each row is the gradient of one output; each column holds the partials with respect to one input. Note the shape convention: $m$ rows, $n$ columns — outputs first.

Two products with the Jacobian matter:

• VJP (reverse mode) — given an upstream row vector $v^\top \in \mathbb{R}^m$, compute $v^\top J \in \mathbb{R}^n$. This is one backward pass, and it never forms $J$.
• JVP (forward mode) — given a direction $u \in \mathbb{R}^n$, compute $Ju \in \mathbb{R}^m$. Cheap when $n$ is small.

A layer that maps $\mathbb{R}^n \to \mathbb{R}^n$ element-wise, like ReLU, has a diagonal Jacobian, so the VJP is just an element-wise multiply — which is why activation layers are cheap.

Worked example

Let $f(x_1, x_2) = (x_1^2 x_2,\; x_1 + x_2)$.

At $(x_1, x_2) = (2, 3)$: $J = \begin{pmatrix} 12 & 4 \\ 1 & 1 \end{pmatrix}$.

• VJP with $v = (1, 0)$, i.e. only $f_1$ matters: $v^\top J = (12, 4)$, the first row.
• VJP with $v = (0, 1)$: $(1, 1)$, the second row.
• VJP with $v = (1, 1)$: $(13, 5)$, the sum of the rows — gradients from both outputs accumulate.

That last case is the general rule: a downstream vector selects a weighted combination of Jacobian rows.

In code

import torch
from torch.func import jacrev

def f(x):
    return torch.stack([x[0]**2 * x[1], x[0] + x[1]])

x = torch.tensor([2.0, 3.0])
print(jacrev(f)(x))  # tensor([[12., 4.], [1., 1.]])

jacrev forms the full matrix for teaching purposes. In real training nobody does: reverse mode computes $v^\top J$ for the single upstream vector $v = \partial L/\partial f$ and skips the $m - 1$ unused rows.

Check yourself

1. What is the shape of the Jacobian of a function $\mathbb{R}^5 \to \mathbb{R}^3$?
2. Under what condition is the Jacobian diagonal, and what does that imply about a VJP?
3. Why is computing a single VJP far cheaper than materializing the Jacobian when $m$ is large?

Key takeaways

• The Jacobian stacks the gradients of every output of a vector function.
• Reverse mode computes $v^\top J$ without ever forming $J$.
• Diagonal Jacobians (element-wise layers) make activation gradients cheap.