When do you write a custom torch.autograd.Function, and how?

## The default is: don't Autograd already differentiates any composition of PyTorch operations, including control flow, indexing and broadcasting. Writing your own `Function` means taking responsibility for correctness that the framework was giving you for free. Assume you don't need one until one of these is true: 1. **You wrapped an external kernel.** A custom CUDA extension, a C++ op or a third-party library call is opaque to autograd, so somebody has to supply the derivative. 2. **The forward is non-differentiable and you want a useful surrogate.** Rounding, sign, argmax, hard thresholding, and quantization all have zero or undefined gradient almost everywhere. The straight-through estimator — forward quantizes, backward passes the gradient through unchanged — exists precisely because you can decouple the two directions in a `Function`. 3. **The composed backward is unstable or slow.** A closed-form derivative you derived by hand can avoid catastrophic cancellation or fuse several kernels into one. 4. **You are trading memory for compute.** A `Function` that recomputes rather than saves is the essence of activation checkpointing; `torch.utils.checkpoint.checkpoint` is the packaged version of that idea, and is what you should use before writing your own. ## The skeleton ``` class Square(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 y = Square.apply(x) ``` The rules that matter: - **Both methods are `@staticmethod`.** There is no `self`; the per-call state lives on `ctx`. - **Call `.apply()`.** Calling `forward` directly runs the numerics but creates no graph node, so the operation silently vanishes from the backward pass. - **`backward` returns one value per `forward` input**, in order. Non-tensor inputs (a flag, an int) get `None`. Returning the wrong arity raises at backward time, not at definition time. - **`grad_output` matches the shape of the forward output.** If forward broadcast or reduced, backward must un-broadcast or expand to match the input shape — a very common source of shape errors. - **`ctx.needs_input_grad`** is a tuple of booleans telling you which inputs actually require gradients; skipping the unneeded computations is a cheap optimisation. ## save_for_backward, and why not just ctx.x = x Stashing a tensor as a plain attribute works numerically but bypasses autograd's bookkeeping. `ctx.save_for_backward` registers the tensor with the version counter, so if it is modified in place before backward runs you get the explicit `RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation` instead of a silently wrong gradient. It also participates in the graph's lifetime management, letting the saved tensor be released when the graph is freed. Use `save_for_backward` for tensors that are inputs or outputs of the forward; plain attributes are fine for small non-tensor metadata such as a shape tuple or a scalar hyperparameter. ## Verifying it A hand-written backward that is subtly wrong does not crash — the model just learns worse, and you will blame the learning rate. So test it: ``` a = torch.randn(4, dtype=torch.double, requires_grad=True) torch.autograd.gradcheck(Square.apply, (a,)) ``` `gradcheck` compares your analytic gradient against a numerical finite-difference estimate and raises on mismatch. Use `dtype=torch.double`: in float32 the finite differences are dominated by rounding noise and the check produces false failures. `torch.autograd.gradgradcheck` does the same for second derivatives, which matters if your `Function` must support `create_graph=True`. For an intentionally non-exact gradient like a straight-through estimator, `gradcheck` will fail by design — that is the point of the estimator — so test those against a hand-computed expectation instead. ## Extras worth knowing - **`setup_context`.** Instead of putting `ctx` in `forward`'s signature, you may define `forward(*inputs)` plus a separate `setup_context(ctx, inputs, output)` static method. That separation is required if you want the `Function` to work with `torch.func` transforms such as vmap. The `ctx`-in-forward spelling remains supported for ordinary use. - **`@torch.autograd.function.once_differentiable`** decorates a `backward` that is not itself differentiable, giving a clear error rather than a confusing one if someone asks for second-order gradients. - **`mark_non_differentiable`** tells autograd that a particular output carries no gradient — useful when a Function returns both a value and an index tensor. - **Export and compile.** A custom `Function` is a black box to tracing and graph capture. If the model must go through `torch.export` or `torch.compile`, prefer registering a proper custom operator so the exporter knows about it, rather than assuming an autograd `Function` will survive the trip.