A backward you can defend

x.norm() and torch.sqrt((x * x).sum()) compute the same number for every input. At x = 0 their gradients differ: the first returns zeros, the second returns NaN. Neither is a bug. The norm has no derivative at the origin, PyTorch's linalg_vector_norm derivative formula special-cases it to zero, and the spelled-out version inherits the derivative of sqrt, which is unbounded as its argument goes to zero.

That gap is the practical form of the chapter's claim that the automatic rule is sometimes the wrong one. Autograd differentiates the ops you called, correctly and one at a time. It has no view of the function you meant, so two decompositions of one function can have two different numerical characters and it will faithfully give you whichever you wrote.

A NaN gradient at the origin is not hypothetical either. Normalizing a vector that can be zero, a distance between two points that can coincide, a standard deviation over a constant window: all of them reach that point during training, and once one NaN enters a parameter, every subsequent step carries it.

import torch

x = torch.zeros(3, requires_grad=True)
x.norm().backward()
print(x.grad)                     # tensor([0., 0., 0.])

y = torch.zeros(3, requires_grad=True)
torch.sqrt((y * y).sum()).backward()
print(y.grad)                     # tensor([nan, nan, nan])

A Function subclass adds one node to the graph in place of everything its forward did. forward computes the value and stashes what backward will need with ctx.save_for_backward, which routes through the same SavedVariable machinery lesson three took apart, version counter included. backward receives one incoming gradient per output and returns one gradient per forward argument.

Two details in the code below are the ones people get wrong first. The return has two entries because forward took two arguments, and the second is None because eps is a float that no gradient can flow to. And ctx.needs_input_grad is a tuple of booleans in the same positions, so a backward that skips work for arguments nobody asked about is a one-line guard rather than a redesign.

The eps inside the square root is the whole fix. Smoothing the norm to sqrt(sum(x*x) + eps*eps) makes the function differentiable everywhere and moves the gradient at the origin to a clean zero, at the cost of a value that is off by eps. Whether that trade is acceptable is a modelling question, and it is one you can only make deliberately once the derivative is yours.

import torch
from torch.autograd import Function, gradcheck, gradgradcheck

class SafeNorm(Function):
    @staticmethod
    def forward(ctx, x, eps):
        n = torch.sqrt((x * x).sum() + eps * eps)
        ctx.save_for_backward(x, n)
        return n

    @staticmethod
    def backward(ctx, g):
        x, n = ctx.saved_tensors
        gx = g * x / n if ctx.needs_input_grad[0] else None
        return gx, None          # one return per forward argument; eps gets None

z = torch.zeros(3, dtype=torch.double, requires_grad=True)
SafeNorm.apply(z, 1e-6).backward()
print(z.grad)                    # tensor([0., 0., 0.], dtype=torch.float64)

t = torch.randn(4, dtype=torch.double, requires_grad=True)
print(gradcheck(lambda a: SafeNorm.apply(a, 1e-6), (t,)))      # True
print(gradgradcheck(lambda a: SafeNorm.apply(a, 1e-6), (t,)))  # True

gradcheck perturbs each input entry by eps, defaulting to 1e-6, and divides the change in each output by it to build a numerical Jacobian. Then it builds the analytic Jacobian by calling your backward once per output entry. If the two agree within atol of 1e-5 and rtol of 1e-3, it returns True; otherwise it raises GradcheckError and prints both matrices.

Those three constants explain the double-precision rule that every tutorial states and few justify. A float32 tensor holds about seven decimal digits, and a perturbation of 1e-6 against values near 1 destroys most of them before the subtraction. The run below is the same correct Square backward checked twice: float32 fails with a numerical column of 2.9802 against an analytic 3.0820, float64 passes.

Read the failure output as a diagnosis rather than a verdict. A numerical column near zero where the analytic one is not means the function is locally flat. Disagreement in the last digits means precision. Disagreement in sign or magnitude means the derivative is wrong. Only the last one is your bug.

import torch
from torch.autograd import Function, gradcheck

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

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

torch.manual_seed(0)
f = torch.randn(2, requires_grad=True)          # float32
try:
    gradcheck(Square.apply, (f,))
except Exception as e:
    print(e)
# Jacobian mismatch for output 0 with respect to input 0,
# numerical:tensor([[ 2.9802,  0.0000],
#         [ 0.0000, -0.5960]])
# analytical:tensor([[ 3.0820,  0.0000],
#         [-0.0000, -0.5869]])
print(gradcheck(Square.apply, (f.detach().double().requires_grad_(True),)))  # True

Rounding has a derivative of zero almost everywhere, and autograd reports exactly that: x.round().sum().backward() fills x.grad with zeros. A quantized layer built on it trains nowhere, because no gradient survives the round.

The standard answer is to lie in one specific place. A straight-through estimator rounds in forward and passes the incoming gradient through untouched in backward, so the rest of the network trains as if the round were the identity. Ten lines, and gradcheck rejects it flatly: the numerical Jacobian is all zeros, the analytic one is the identity, and the error prints both.

That failure is the correct result and the reason to run the check anyway. Now the disagreement is documented instead of assumed, you know it is total rather than partial, and the next person to read the class can see which of the two Jacobians was chosen on purpose. A custom Function is a claim about a derivative, and running the check turns the claim into either a green result or a known, deliberate exception.

import torch
from torch.autograd import Function, gradcheck

class RoundSTE(Function):
    @staticmethod
    def forward(ctx, x):
        return torch.round(x)

    @staticmethod
    def backward(ctx, g):
        return g

x = torch.tensor([0.3, 1.7], requires_grad=True)
x.round().sum().backward()
print(x.grad)                    # tensor([0., 0.]): the true derivative

y = torch.tensor([0.3, 1.7], requires_grad=True)
RoundSTE.apply(y).sum().backward()
print(y.grad)                    # tensor([1., 1.]): the one that trains

t = torch.tensor([0.3, 1.7], dtype=torch.double, requires_grad=True)
try:
    gradcheck(RoundSTE.apply, (t,))
except Exception as e:
    print(type(e).__name__)      # GradcheckError
    print(e)
# Jacobian mismatch for output 0 with respect to input 0,
# numerical:tensor([[0., 0.],
#         [0., 0.]], dtype=torch.float64)
# analytical:tensor([[1., 0.],
#         [0., 1.]], dtype=torch.float64)

Three of the four decisions in SafeNorm are not PyTorch decisions. Something has to hold the values backward will read, something has to know which inputs want gradients, and something has to check the analytic rule against a numerical one. The JAX path names those res in a custom_vjp residual tuple, the argnums of the transform, and jax.test_util.check_grads.

What is genuinely local to this arc is where the state lives. A ctx object is mutable, per call, and carries a version counter into a comparison that happens later; a residual tuple is a value returned from one function and passed to another. That is the same split the modules chapter draws between state on an object and state threaded through a signature, showing up one layer down in the autodiff machinery itself.

Everything else in this arc holds for both. A tape is nodes and edges, a node keeps what its rule needs, and a rule you wrote yourself is worth exactly as much as the check you ran against it.