PyTorch Jacobians with torch.func
A Jacobian is a table that holds every output gradient for every input. torch.func builds that table in one call on torch 2.14.1+cpu.
One backward per output row
Plain autograd fills one row per backward call here. Two outputs mean two calls, and the first call must keep its graph.
import torch
def f(x):
return torch.stack([x[0] ** 2 + x[1] * x[2], 3 * x[0] + x[2] ** 2])
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
out = f(x)
row0 = torch.autograd.grad(out[0], x, retain_graph=True)[0]
row1 = torch.autograd.grad(out[1], x)[0]
print("Row for output 0:", row0.tolist())
print("Row for output 1:", row1.tolist())jacrev builds the whole table
jacrev runs the backward style of differentiation once per output. The table prints with one row per output and one column per input.
from torch.func import jacrev
J = jacrev(f)(torch.tensor([1.0, 2.0, 3.0]))
print(J)
print("Shape:", tuple(J.shape))jacfwd walks the inputs instead
jacfwd pushes one input direction at a time through the function. Its table matches the jacrev table exactly on this function.
from torch.func import jacfwd
J2 = jacfwd(f)(torch.tensor([1.0, 2.0, 3.0]))
print("Same matrix:", torch.equal(J, J2))Example
import torch
from torch.func import jacrev, jacfwd
def f(x):
return torch.stack([x[0] ** 2 + x[1] * x[2], 3 * x[0] + x[2] ** 2])
x = torch.tensor([1.0, 2.0, 3.0])
print("Input:", x.tolist())
print("Output:", f(x).tolist())
print("Jacobian with jacrev:")
print(jacrev(f)(x))
print("Jacobian with jacfwd:")
print(jacfwd(f)(x))
hand = torch.tensor([[2.0, 3.0, 2.0], [3.0, 0.0, 6.0]])
print("Hand-checked Jacobian:")
print(hand)
print("jacrev matches hand:", torch.equal(jacrev(f)(x), hand))
print("jacfwd matches hand:", torch.equal(jacfwd(f)(x), hand))Key points
- jacrev: builds the Jacobian with reverse mode, one output row at a time
- jacfwd: builds the Jacobian with forward mode, one input column at a time
- Use jacfwd when the function has fewer inputs than outputs
- Use jacrev when the function has fewer outputs than inputs
- Both modes return the same table, and only the cost differs
- In short, the Jacobian packs all first derivatives of a vector function into one table
In this example, f maps three inputs to two outputs. The first output adds x0 squared to the product of x1 and x2. The second output adds three times x0 to x2 squared. At the input 1, 2, and 3, the outputs print as 7 and 12. jacrev and jacfwd each build the full two by three table in one call. By hand, the first row is 2, 3, and 2 from differentiating against each input in turn. The second row is 3, 0, and 6 by the same working. Both printed tables match the hand table, and the two equality checks print True.
Note
torch.func is version-sensitive, and these lessons pin torch 2.14.1+cpu
Try it Yourself
Edit the code, then press Run.
Running is simulated · Python Playground
Exercise
A function takes five inputs and returns two outputs. Which builder walks fewer passes?