torch.func — Per-Sample Gradients and Vmap

Functional autograd, batched transforms

torch.func (the in-tree replacement for the old standalone functorch) gives you JAX-style function transformations: grad, vmap, jacrev, hessian. They let you write code that operates on a single example and then automatically vectorize it over a batch.

Two practical use cases that come up surprisingly often:

Per-sample gradients. Standard .backward() gives you the gradient of the summed loss across the batch — every parameter has one gradient. Some research (differential privacy, influence functions, GradSAM) needs the gradient of every individual sample's loss. torch.func.vmap(grad(...)) gives you that without writing a Python loop.
Higher-order gradients. Hessian-vector products, second-order optimizers, meta-learning — all need gradients of gradients. torch.func.grad(grad(f)) composes cleanly.

The mental shift

Standard PyTorch: tensors carry implicit graph state, you call .backward(). torch.func: pure functions of inputs and parameters, transforms produce new pure functions. It's a small but real mindset change — closer to JAX, slightly less ergonomic for normal training, much more powerful for the cases above.

Code

grad as a function transform·python

import torch
from torch.func import grad

def f(x):
    return torch.sin(x) * x

# grad(f) is a NEW function: x → df/dx
df_dx = grad(f)

x = torch.tensor(1.0)
print(df_dx(x))            # cos(1)*1 + sin(1) ≈ 0.541 + 0.841 = 1.381

# Higher order — grad(grad(f))
d2f_dx2 = grad(grad(f))
print(d2f_dx2(x))          # -sin(1)*1 + 2*cos(1) ≈ ...

vmap — auto-batch over a single-example function·python

import torch
from torch.func import vmap, grad

def loss_per_example(w, x, y):
    pred = (w * x).sum()
    return (pred - y) ** 2

w = torch.randn(4)
batch_x = torch.randn(8, 4)         # batch of 8
batch_y = torch.randn(8)

# Single-sample gradient: grad w.r.t. w
single_grad = grad(loss_per_example)

# Vectorize over the batch dim of x and y (in_dims=(None, 0, 0))
per_sample_grads = vmap(single_grad, in_dims=(None, 0, 0))(w, batch_x, batch_y)
print(per_sample_grads.shape)  # torch.Size([8, 4]) — one gradient per sample

Per-sample gradients on an nn.Module — the modern way·python

import torch
import torch.nn as nn
from torch.func import functional_call, vmap, grad

model = nn.Linear(4, 2)
params = dict(model.named_parameters())

def compute_loss(params, x, y):
    pred = functional_call(model, params, (x.unsqueeze(0),)).squeeze(0)
    return ((pred - y) ** 2).mean()

batch_x = torch.randn(16, 4)
batch_y = torch.randn(16, 2)

# Per-sample gradient w.r.t. params
per_sample_grad = vmap(grad(compute_loss), in_dims=(None, 0, 0))(
    params, batch_x, batch_y
)
print({k: v.shape for k, v in per_sample_grad.items()})
# {'weight': torch.Size([16, 2, 4]), 'bias': torch.Size([16, 2])}

torch.func — Per-Sample Gradients and Vmap

Functional autograd, batched transforms

The mental shift

Code

External links

Exercise

Progress

Comments 0