Custom autograd Functions

Needed for straight-through estimators, numerically stable gradients, and wrapping non-PyTorch code.

When autograd's automatic derivative is not what you want, subclass torch.autograd.Function and write both directions yourself:

forward(ctx, *inputs): compute the output. Store what backward needs with ctx.save_for_backward(*tensors); non-tensor values can be set as attributes on ctx.
backward(ctx, *grad_outputs): receive the gradient of the loss with respect to each output and return one gradient per forward input, or None for inputs that need none (like integer arguments).

Both methods are @staticmethods, and you call the function with MyFn.apply(x), never by instantiating it. Recent PyTorch versions also support splitting out a setup_context(ctx, inputs, output) method, which is required for compatibility with torch.func transforms.

Typical uses:

Straight-through estimators: round or binarise in the forward pass but pretend it was identity in the backward pass, as in quantisation-aware training.
Stable or cheaper gradients for a fused formula.
Wrapping code autograd cannot see, such as a NumPy or C++ routine.

Always check a hand-written backward with torch.autograd.gradcheck, which compares it to finite differences. Use float64 inputs, because float32 is too imprecise for the check.

Example

import torch
from torch.autograd import Function
class MyExp(Function):
    @staticmethod
    def forward(ctx, x):
        y = x.exp()
        ctx.save_for_backward(y)        # exp's derivative is exp itself
        return y
    @staticmethod
    def backward(ctx, grad_out):
        (y,) = ctx.saved_tensors
        return grad_out * y
class RoundSTE(Function):
    @staticmethod
    def forward(ctx, x):
        return x.round()
    @staticmethod
    def backward(ctx, grad_out):
        return grad_out                 # pretend round() was identity

x = torch.randn(3, dtype=torch.double, requires_grad=True)
print(torch.autograd.gradcheck(MyExp.apply, (x,)))   # True
z = torch.tensor([0.4, 1.6], requires_grad=True)
RoundSTE.apply(z).sum().backward()
print(z.grad)                           # tensor([1., 1.])  (true derivative is 0)

Remember

Static forward/backward; call with Fn.apply(...)
Return one gradient per input, None where not needed
Verify with gradcheck on float64

Watch out for

Returning fewer gradients from backward than forward had inputs, which raises an error.

Check yourself

forward(ctx, x, k) takes a tensor x and a Python int k. What should backward return?

Official docs: pytorch.org/docs/stable/notes/extending.html

Example

Remember

Watch out for

Check yourself

Builds on

Unlocks