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).