Your own backward pass

Autograd differentiates by composing the derivatives of primitives. When you need to replace one of those, you write a torch.autograd.Function.

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

    @staticmethod
    def backward(ctx, grad_out):
        (x,) = ctx.saved_tensors
        return grad_out * 2 * x

Two static methods and a context object between them.

The contract

• forward computes the value and saves whatever backward will need, via ctx.save_for_backward. Saving with a plain attribute works for non-tensors and skips the version checking that catches a saved tensor being mutated afterwards.
• backward receives grad_out, which is dL/d(output), and must return dL/d(input) for each input to forward, in the same order. Return None for inputs that do not need a gradient.
• backward runs with grad disabled, so operations inside it do not build a graph of their own. If you need a second derivative, that is what torch.autograd.grad(create_graph=True) is for.

Call it with Square.apply(x), not Square()(x).

Why bother

Three real reasons, and “it is faster” is rarely one of them:

• Numerical stability. The composed derivative of a formula can overflow where the analytic one does not. This is why logsumexp and cross_entropy are single primitives rather than compositions.
• Memory. Recomputing an input in backward instead of saving it trades time for memory, which is exactly what gradient checkpointing does.
• Non-differentiable steps. A straight-through estimator passes the gradient around a quantisation or a rounding that has no useful derivative.

Checking it

torch.autograd.gradcheck(Square.apply, (x.double().requires_grad_(),))

compares your backward against a numerical estimate. Use float64: in float32 the finite differences are noisier than the thing being measured, and gradcheck will fail a correct implementation.

Your task

Implement Cube, an autograd Function computing x ** 3 with a correct backward, then:

def cube_grad(x: torch.Tensor) -> torch.Tensor

returning the gradient of Cube.apply(x).sum() with respect to x.

The starter’s backward forgets the chain rule.