🎓 A1Learner
🔍

PyTorch Custom Autograd Functions

A custom autograd function is a class that supplies its own forward and backward rules. Autograd calls the backward rule during backward() and trusts the gradients it returns.

The built-in clamp keeps its own rule

The built-in clamp caps the value at 1.0 here. Its fixed backward rule passes no gradient outside the range.

import torch

x = torch.tensor([2.0], requires_grad=True)
y = x.clamp(-1.0, 1.0)
print("Clamped value:", y.item())
y.backward()
print("Gradient outside the range:", x.grad.item())

A first custom function

DoubleFn doubles its input in forward. Its backward doubles the incoming gradient, because the derivative of two times x is two.

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

    @staticmethod
    def backward(ctx, grad_output):
        return grad_output * 2


x = torch.tensor([3.0], requires_grad=True)
y = DoubleFn.apply(x)
print("Forward:", y.tolist())
y.backward()
print("Gradient:", x.grad.tolist())

Save the input for backward

The square rule needs the input again inside backward. save_for_backward stores it, and backward reads it back from ctx.

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


x = torch.tensor([4.0], requires_grad=True)
y = SquareFn.apply(x)
y.backward()
print("Forward:", y.item())
print("Gradient:", x.grad.item())

Example

import torch


class ClampFn(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, low, high):
        ctx.save_for_backward(x)
        ctx.low = low
        ctx.high = high
        return x.clamp(low, high)

    @staticmethod
    def backward(ctx, grad_output):
        (x,) = ctx.saved_tensors
        inside = (x >= ctx.low) & (x <= ctx.high)
        return grad_output * inside.to(grad_output.dtype), None, None


x = torch.tensor([-2.0, -0.5, 0.3, 2.0], requires_grad=True)
y = ClampFn.apply(x, -1.0, 1.0)
print("Input:", [round(v, 2) for v in x.detach().tolist()])
print("Clamped:", [round(v, 2) for v in y.tolist()])
y.sum().backward()
print("Gradient:", x.grad.tolist())

Key points

  • forward: the method that computes the output from the inputs
  • backward: the method that turns an incoming gradient into gradients for the inputs
  • ctx: the context object that carries saved values from forward to backward
  • Use save_for_backward on ctx to keep tensors that backward needs
  • Call a custom function through its apply method, never by calling the class
  • In short, forward does the maths and backward does the derivative, both written by hand

In this example, ClampFn is a custom autograd function with two saved pieces of state. Its forward clamps the input into the range from minus one to one. It saves the input on ctx and stores the two bounds beside it. Its backward builds a mask of the inputs inside the range. It returns the incoming gradient times that mask, plus None for each bound. The input row holds minus two, minus a half, 0.3, and two. The clamped row prints the two outside values pulled to the bounds. After sum and backward, the gradient prints one inside the range and zero outside. The two None returns match the two bound arguments, which need no gradient.

Note

A wrong backward still runs, and only a numerical check can expose it

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What must backward return for an input that needs no gradient, like a bound?