Custom autograd functions

Custom Functions are where the VJP becomes your code. One correct handwritten backward removes the mystery from every built-in one.

derive
implement

You can now

Write forward, setup_context and backward for a Function subclass, with save_for_backward and None grads for non-tensor args.
Validate the backward with gradcheck in float64, and state whether double backward is supported.
Explain why saved state must go through save_for_backward rather than raw ctx attributes.

Covers

Custom autograd.Function forward and backward

Study, in this order

2 steps

Extending `torch.autograd`: custom `Function` contract and examples
PyTorch - Autograd mechanics, serialization and reproducibility

When to use and When not to use; How to use steps 1 to 4 (forward, setup_context, backward, ctx rules, once_differentiable, gradcheck); Example (LinearFunction, MulConstant, gradcheck with dtype=torch.double, eps=1e-6, atol=1e-4); Combined or separate forward() and setup_context()

Depth · Skip the __torch_function__ protocol, C++ extensions, forward-mode jvp and torch.func transforms. Stop at Combined or separate forward() and setup_context().

Done
Autograd mechanics: how saved tensors work
PyTorch - Autograd mechanics, serialization and reproducibility

Saved tensors (what pow and exp save, packing versus unpacking, saved-tensor hooks)

Depth · Read only the Saved tensors section here; the grad-mode and in-place sections belong to earlier units.

Done

Focus on

forward() runs with grad recording off, so nothing inside it is tracked. setup_context() owns the ctx writes. backward() receives one grad per output and returns one grad per forward input.
Save tensors with save_for_backward. The Extending note gives the reasons: the engine frees them when this backward finishes, it avoids reference cycles when you save an output, and saved-tensor hooks (activation checkpointing, offloading) can see them. Saved tensors are also version-checked, so a later in-place edit raises an error instead of corrupting the gradient.
Non-tensor forward args get None grads, and backward must never modify its grad inputs in place. Gate optional work with needs_input_grad.
Derive the linear backward shapes on paper first: grad_input = grad_output mm weight, grad_weight = grad_output.t() mm input, grad_bias = grad_output.sum(0).

Reproduce closed-book

From a blank file, reimplement LinearFunction and MulConstant with separate setup_context, and gradcheck both in float64. Then break the bias grad (return zeros) and quote the gradcheck mismatch before fixing it.

Traps

Stashing tensors as plain ctx attributes. They stay alive for the whole graph, can form a reference cycle with the output, escape the in-place version check, and are invisible to saved-tensor hooks. Keep ctx attributes for non-tensors.
Writing into grad_output inside backward. Grad inputs are read-only; in-place edits corrupt every other consumer of that graph.
Returning a tensor grad for a non-tensor arg like constant. Those positions must be None, as MulConstant shows.

Practice

Extending PyTorch worked examples: `LinearFunction` and `MulConstant`, each gradchecked

Example: reimplement LinearFunction and both MulConstant variants from a blank file, then run the gradcheck snippet (double inputs, eps=1e-6, atol=1e-4) on each

Do not copy the page. Write from the contract (forward, setup_context, backward), pass gradcheck, then deliberately break one grad and read the failure.

Self-test

Linear forward is input.mm(weight.t()) with input (5, 4) and weight (3, 4). What are the shapes of grad_input and grad_weight for grad_output (5, 3)?
Why does MulConstant backward return two values with the second None?
With separate forward and setup_context, you want to save an intermediate computed in forward that is neither an input nor an output. What do you do?

Completion criterion

Stop when

Both reimplementations pass gradcheck from memory, and you can state for each saved value why it is saved and when it is freed.

Mark passed, then open 2.6 Only mark this when you can do it without notes.

Leads to

8.1 Eager execution, autograd nodes and saved tensors See what saved tensors cost at runtime.
1.4 Branches, broadcasting and saved values in backward Rebuild the same VJP logic in your own engine.

Comes after

2.4 Backward, `zero_grad`, `gradcheck` and evaluation modesGradcheck the backward you write yourself.
8.1 Eager execution, autograd nodes and saved tensorsCustom Functions build directly on save_for_backward and version rules.

Resources for this unit

PyTorch - Autograd mechanics, serialization and reproducibility
PyTorchrequiredreference
Read Autograd mechanics, Serialization semantics and Reproducibility. Test leaf behavior, saved tensors, grad modes, state dictionaries and the exact tolerance promised by one interrupted and resumed run.

Also in this section

PyTorch - Learn the Basics
PyTorchprimarycourse
Complete Tensors, Datasets and DataLoaders, Transforms, Build Model, Autograd, Optimization, and Save and Load Model. Rebuild each essential example in a small src package instead of retaining a tutorial notebook.
Learn PyTorch - selected framework exercises
Daniel Bourkerequiredpractice
Complete exercise 00, tasks 1-5 in 01, tasks 1-7 in 02, tasks 5-13 in 03, tasks 1-3 in 05, and exercises 1-2 in 07. Recreate solutions independently and verify behavior against current official docs.
Deep Learning with PyTorch, 2nd edition - framework foundations
Howard Huang, Luca Antiga, Eli Stevens and Thomas Viehmann (Manning)supportimplementation
Read Sections 3.8-3.9 and 5.5 when tensor storage, strides or autograd need a second explanation, with the Chapter 3 and 5 companion notebooks for runnable cells. Inspect the Chapter 4 and 7 code only for a diagnosed data-representation or loading gap.
NeetCode ML - PyTorch closing drills
NeetCodeoptionalpractice
After the module gate, solve Pytorch Basics, Layer Normalization, Batch Normalization and RMS Normalization cold. Explain shapes, state and train-evaluation behavior aloud. These do not replace the package and resume test.
PyTorch Automatic Mixed Precision recipe - Saving/Resuming
PyTorchreferencereference
Read the Saving/Resuming section: save scaler.state_dict() with the model and optimizer state, at the start of an iteration or after scaler.update(), and how to resume with or without AMP.
Hands-On Machine Learning with Scikit-Learn and PyTorch
Aurelien Geron (O'Reilly, 2025)supporttheory
Read Chapter 10, Implementing Mini-Batch Gradient Descent Using DataLoaders and Model Evaluation, for the train and evaluate loops and how batch metrics must be aggregated.