Skip to content

← Autograd Mechanics step 4 of 8

Hard Primitives

Your own backward pass

Jump to the task

A custom torch.autograd.Function defines a forward calculation and its backward rule. For an elementwise square, the rule multiplies the incoming gradient by the local derivative:

import torch

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

x = torch.tensor([2.], requires_grad=True)
loss = (3 * Square.apply(x)).sum()
loss.backward()
print(x.grad.tolist())  # [12.0]

The local derivative at x=2 is four. The downstream multiplication sends an incoming gradient of three, so the final gradient is twelve. Call the function through Square.apply, rather than instantiating it.

Save values and return input gradients

Use ctx.save_for_backward for tensors the backward calculation needs. This supports saved-tensor handling, including version checks. Non-tensor metadata can be stored as context attributes.

Backward receives a gradient for each output and returns a gradient for each forward input, in input order. Use None where no gradient is supplied, such as for a non-tensor argument. For this one-input elementwise function, backward returns grad_out * local_derivative.

Ordinary backward execution does not record a new gradient graph. With create_graph=True, backward operations can themselves be recorded. A custom backward must still be written to support the higher derivatives you need; the flag alone does not fix an incompatible implementation.

When a custom function is useful

Prefer ordinary differentiable PyTorch operations when they express your calculation. Custom functions can connect an external implementation to autograd, implement a specialized derivative, or trade saved values for recomputation. A straight-through estimator deliberately supplies a surrogate gradient; it is not the true derivative of the discontinuous forward rule.

Check the derivative

x = torch.tensor([0.3, -0.7], dtype=torch.float64, requires_grad=True)
print(torch.autograd.gradcheck(Square.apply, (x,)))  # True

Gradcheck compares the backward rule with finite-difference estimates. Double precision is the usual choice for its default tolerances. Lower precision, nondifferentiable points or unsuitable step sizes can make the numerical comparison unreliable. Passing a few checks is useful evidence, not proof for every input or derivative order.

Your task

Implement Cube so its forward returns x ** 3 and its backward applies the chain rule. Then implement:

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

Compute y = Cube.apply(x) and return the gradient of (2 * y).sum() with respect to x. The harness supplies a fresh float64 leaf with gradients enabled and x.grad initially None. The expected gradient is 6 * x ** 2.

The factor two is deliberate: a plain sum sends all-ones incoming gradients and would not expose the starter's omission of grad_out. Preserve the weighted loss while fixing the custom backward.