Custom Autograd.Function¶

Overview¶

torch.autograd.Function is the escape hatch for operations that autograd doesn't know how to differentiate— either because they're custom CUDA kernels, non-differentiable intermediates with hand-derived gradients, or you want to fuse ops so backward uses fused kernels too.

Subclass Function; implement static forward(ctx, *inputs) and backward(ctx, *grad_outputs).
ctx.save_for_backward(...) saves tensors for the backward pass (cheaper than closures, and respected by checkpoints).
apply builds the graph node: call MyFn.apply(x)— not MyFn()(x).
backward must return the same count of gradients as forward received inputs (or None).

Materialized saved tensors dominate memory in transformers— custom fns re-compute intermediates to trade compute for memory.

-

Minimal Example— Differentiable "Square"¶

import torch

class SquareFn(torch.autograd.Function):
 @staticmethod
 def forward(ctx, x):
 ctx.save_for_backward(x)
 return x * x

 @staticmethod
 def backward(ctx, grad_output):
 (x,) = ctx.saved_tensors
 return grad_output * 2 * x # d(x^2)/dx = 2x

x = torch.tensor([2.0, 3.0], requires_grad=True)
y = SquareFn.apply(x)
y.sum().backward()
print(x.grad) # tensor([4., 6.])

Why `apply` and not `call`?¶

apply registers the op in the autograd graph; plain Python call on the instance wouldn't build nodes.

-

A Realistic One— Clipped ReLU (with dead zone)¶

class ClippedReLU(torch.autograd.Function):
 @staticmethod
 def forward(ctx, x, max_val):
 ctx.save_for_backward(x)
 ctx.max_val = max_val
 return torch.clamp(x, min=0, max=max_val)

 @staticmethod
 def backward(ctx, grad_output):
 (x,) = ctx.saved_tensors
 mask = ((x > 0) & (x < ctx.max_val)).float()
 return grad_output * mask, None # None -> max_val is not a tensor

x = torch.tensor([-1.0, 0.5, 3.0], requires_grad=True)
out = ClippedReLU.apply(x, 2.0)
out.sum().backward()
print(x.grad) # [0., 1., 0.]

Return None for every non-tensor/statically-constant forward arg.

-

Fused Ops— Save Memory by Recomputing¶

class MatmulFn(torch.autograd.Function):
 """y = a @ b; backward recomputes nothing extra (uses saved a,b)."""
 @staticmethod
 def forward(ctx, a, b):
 ctx.save_for_backward(a, b)
 return a @ b

 @staticmethod
 def backward(ctx, g):
 a, b = ctx.saved_tensors
 return g @ b.t(), a.t() @ g # dL/da, dL/db

a = torch.randn(4, 8, requires_grad=True)
b = torch.randn(8, 3, requires_grad=True)
MatmulFn.apply(a, b).sum().backward()
print("a.grad shape:", a.grad.shape)

-

`double_backward`— Arbitrary High-Order¶

class CubeFn(torch.autograd.Function):
 @staticmethod
 def forward(ctx, x):
 ctx.save_for_backward(x)
 return x ** 3

 @staticmethod
 def backward(ctx, g):
 (x,) = ctx.saved_tensors
 return g * 3 * x * x

 # d/dx (g * 3x^2) needed if create_graph=True used:
 @staticmethod
 def double_backward(ctx, gg):
 (x,) = ctx.saved_tensors
 return gg * 6 * x

x = torch.tensor([2.0], requires_grad=True)
y = CubeFn.apply(x)
g1 = torch.autograd.grad(y, x, create_graph=True)[0]
g2 = torch.autograd.grad(g1, x)[0]
print(g1.item(), g2.item()) # 12.0 12.0 (3*2^2, 6*2)

Pitfalls & Rules¶

Pitfall	Fix
Forward returns but backward count mismatched	identical # inputs incl. `None`
Saving tensors that later get mutated in-place	`save_for_backward` + don't mutate saved
Non-tensor args in `forward`	pass through `ctx` as attributes
Forgetting `.apply()`	graph node never created
Relying on re-forward in backward	re-run forward inside `backward` only with `ctx.needs_input_grad` guards

Guarding input-grad necessity¶

class Chained(Function):
 @staticmethod
 def forward(ctx, x):
 r = x * 2
 ctx.save_for_backward(r)
 return r * 3
 @staticmethod
 def backward(ctx, g):
 (r,) = ctx.saved_tensors
 gi = g * 3 * 2 if ctx.needs_input_grad[0] else None
 return gi

-

Key Takeaways¶

Custom Function = explicit forward + backward with ctx-saved tensors.
Use it for fused/custom ops, non-automatic gradients, and memory-vs-compute trade-offs.
Match backward arity exactly; None for non-tensor args.
double_backward handles create_graph=True usage (learned optimizers, meta-learning).
Need a custom elementwise CUDA op? Pair Function with a C++/CUDA extension (Chapter 9).

-

Autograd
[Higher Order Gradients](/06-pytorch/03-custom-autograd-and-model-transformation/(02-higher-order-gradients-jacobians/)
Custom Kernels