torch.func

Active contributors: zou3519, Chillee, kshitij12345

Purpose

torch.func (formerly functorch) is the function-transform package: composable transforms like vmap, grad, jvp, vjp, jacrev, jacfwd, hessian, and functional_call. It is the JAX-style API on top of PyTorch.

This is the user-facing surface of the Functorch system.

Quick examples

import torch
from torch.func import grad, vmap, jacrev, functional_call

# Per-sample gradients (vmap of grad)
def loss_fn(params, x, y):
    out = torch.func.functional_call(model, params, x)
    return ((out - y) ** 2).mean()

grad_per_sample = vmap(grad(loss_fn), in_dims=(None, 0, 0))(params, xs, ys)

# Jacobian
J = jacrev(model)(x)

# Hessian
H = torch.func.hessian(loss_fn)(params)

Public API

Function	What it does
vmap(f, in_dims)	Vectorize f over a batch dim
grad(f)	Reverse-mode gradient wrt the first argument
grad_and_value(f)	Same as grad but also returns the function value
vjp(f, *primals)	Returns (out, vjp_fn) for reverse-mode pullback
jvp(f, primals, tangents)	Forward-mode JVP
jacrev(f) / jacfwd(f)	Per-element jacobian
hessian(f)	Second-order derivative
functional_call(module, params_dict, args)	Run module with externally-provided params
stack_module_state([m1, m2, ...])	Stack params/buffers from N copies of a module for vmap
replace_all_batch_norm_modules_(model)	Replace BN with running-stat-free variant for vmap

Composability

The whole point: transforms compose freely.

torch.func.vmap(torch.func.grad(loss_fn))    # per-sample gradients
torch.func.vmap(torch.func.vmap(f))           # nested vmap → arbitrary batching
torch.func.jacrev(torch.func.jacrev(f))       # Hessian via two reverse-mode passes
torch.func.jvp(torch.func.grad(f), ..., ...)  # mixed forward+reverse

Where it sits in the stack

The transforms work by pushing dispatch keys onto the active keyset. When ATen ops run inside the transform context, the dispatcher routes them to the transform's batch rule / grad wrapper / jvp wrapper kernel. See Functorch for details.

Stateless model calls

The classic PyTorch idiom of model.forward(x) couples weights to the module instance. For higher-order optimization, ensembles, MAML, or anything that wants to call a model with other weights, use functional_call:

# meta-learning step: compute loss with current weights, then with weights one step ahead
loss1 = functional_call(model, params, x)
grads = torch.func.grad(...)(params, x, y)
new_params = {k: v - lr * grads[k] for k, v in params.items()}
loss2 = functional_call(model, new_params, x)

Limitations

• Not all ATen ops have batch rules. vmap falls back to an explicit loop for ops without rules; expect warnings.
• No mutation. Transformed functions must be pure (no .add_, no inplace).
• Dynamic control flow on tensor values. Works in eager but won't trace through torch.compile cleanly.
• torch.compile(torch.vmap(...)) is supported but is more sensitive to op coverage than eager.

Where to look

File	Purpose
torch/_functorch/apis.py	The torch.func re-exports
torch/_functorch/eager_transforms.py	grad / vjp / jvp / jacrev / jacfwd
torch/_functorch/vmap.py	vmap
torch/_functorch/functional_call.py	functional_call
aten/src/ATen/functorch/	C++ side

Where to read next

• Systems / Functorch — implementation.
• Systems / AOT Autograd — production user of these transforms.