Functorch

Active contributors: zou3519, Chillee, kshitij12345

Purpose

functorch (now exposed as torch.func) is the JAX-style function-transform package: composable transforms like vmap, grad, jvp, vjp, jacfwd, jacrev, hessian, and functional_call. It is also the home of AOT Autograd (covered separately in AOT Autograd).

The Python-facing API lives at torch/_functorch/ (with a thin re-export at torch/func/); the C++ side that implements the dispatch keys lives at aten/src/ATen/functorch/ and functorch/.

Directory layout

Path	Contents
torch/_functorch/	Python implementations and AOT Autograd
torch/_functorch/eager_transforms.py	grad, vjp, jvp, jacrev, jacfwd, hessian
torch/_functorch/vmap.py	vmap user-facing API
torch/_functorch/functional_call.py	Stateless module call
torch/_functorch/apis.py	Exposed at torch.func
torch/_functorch/_aot_autograd/	AOT Autograd (used by torch.compile)
torch/_functorch/partitioners.py	fwd/bwd partitioners
torch/_functorch/aot_autograd.py	Top-level entry
aten/src/ATen/functorch/	C++ side of the transforms — dispatch-key kernels
aten/src/ATen/functorch/BatchRulesBinaryOps.cpp	Per-op vmap batch rules
aten/src/ATen/functorch/Interpreter.h	Interpreter stack used by transforms
functorch/	Top-level historical directory; mostly examples and tests now

Key abstractions

Type	File	Purpose
vmap(f, in_dims=...)	torch/_functorch/vmap.py	Vectorize f over a batch dim
grad(f)	torch/_functorch/eager_transforms.py	Reverse-mode gradient as a function transform
jvp(f, primals, tangents)	torch/_functorch/eager_transforms.py	Forward-mode JVP
vjp(f, *primals)	torch/_functorch/eager_transforms.py	Reverse-mode VJP
jacrev / jacfwd	torch/_functorch/eager_transforms.py	Per-element jacobian via vmapped grad / jvp
functional_call	torch/_functorch/functional_call.py	Run a module with externally-provided params
Interpreter	aten/src/ATen/functorch/Interpreter.h	The transform stack the dispatcher pushes through
BatchedTensorImpl	aten/src/ATen/functorch/BatchedTensorImpl.h	Wrapper tensor for vmap
TensorWrapper (Grad)	aten/src/ATen/functorch/TensorWrapper.h	Wrapper tensor for grad/vjp

How it works

Function transforms are implemented as a stack of dispatch-key wrappers, not as Python rewrites. Each transform pushes an Interpreter onto a per-thread stack and a corresponding dispatch key (FuncTorchBatched, FuncTorchGradWrapper, FuncTorchVmapMode) onto the active keyset. When ATen ops are called inside the transform, the dispatcher routes them to the transform's kernel first.

vmap

Inputs are wrapped in BatchedTensorImpls that carry a bdim (the batch dimension index). The FuncTorchBatched kernel for each op (the batch rule) tells the dispatcher what to do:

• For most pointwise ops, the rule moves the batch dim to the front and redispatches to the same op.
• For ops that don't naturally vectorize (e.g., nonzero), the rule loops over the batch dim using LegacyBatchedFallback.

Batch rules live in aten/src/ATen/functorch/BatchRules*.cpp (a few thousand entries). Many are auto-generated; many are hand-written.

graph TB
    User[user fn f] -->|vmap| Wrap[wrap inputs in BatchedTensor]
    Wrap --> Op[run f]
    Op -->|each ATen call| Disp[Dispatcher]
    Disp -->|FuncTorchBatched key| BR[batch rule for that op]
    BR -->|adjust bdim, redispatch| Disp
    Disp -->|backend key| Kernel[CPU/CUDA kernel]

grad / vjp / jvp

grad(f) runs f inside a wrapper that:

1. Pushes a FuncTorchGradWrapper interpreter and key.
2. Wraps inputs in TensorWrappers that record forward outputs.
3. Calls regular autograd to compute backward.
4. Returns the gradients as values.

The wrapper layer lets grad(grad(f)) and vmap(grad(f)) compose: each transform stacks its own interpreter, and the dispatcher walks through them in order. This is the same machinery JAX uses, modelled directly with the dispatcher.

jacrev / jacfwd / hessian

These are implemented as obvious compositions: jacrev(f) = vmap(vjp(f)), jacfwd(f) = vmap(jvp(f)), hessian(f) = jacrev(jacrev(f)). The composability is the point.

functional_call

torch.func.functional_call(module, params_dict, args) calls a module's forward with externally-supplied parameter values, without mutating the module. It walks the module to swap parameters/buffers in, runs forward, and swaps them back. Used for higher-order optimizers (MAML, hypernetworks, BERT-style finetuning) and for tracing modules in transform-friendly form.

Integration points

• Autograd. grad/vjp ride on top of the regular autograd engine. See Autograd.
• Dispatcher. Everything is layered via FuncTorch dispatch keys. See Dispatcher.
• AOT Autograd. Lives in this same package and is the production user of these transforms. See AOT Autograd.
• Compile. torch.compile(torch.vmap(f)) works because the transforms are pushed before Dynamo's tracing starts.

Entry points for modification

• New vmap rule for an op → add a batch rule in aten/src/ATen/functorch/BatchRules*.cpp.
• New transform → add a wrapper module under torch/_functorch/. Most user-facing transforms compose existing ones rather than introducing new dispatch keys.
• For debugging vmap fallbacks, set TORCH_LOGS=batch_rule to see which ops fall back to the slow generic path.

Key source files

File	Purpose
torch/_functorch/vmap.py	vmap user API
torch/_functorch/eager_transforms.py	grad, vjp, jvp, jacrev, jacfwd
torch/_functorch/functional_call.py	Stateless module call
aten/src/ATen/functorch/Interpreter.h	Transform interpreter stack
aten/src/ATen/functorch/BatchedTensorImpl.h	Wrapper tensor for vmap
aten/src/ATen/functorch/BatchRulesBinaryOps.cpp	Per-op batch rules