Custom Autograd & Model Transformation

Overview

Two skills separate advanced PyTorch from copy-paste usage:

1. Custom autograd.Function— defining your own forward/backward pair for ops PyTorch doesn't have (or needs fused/special handling for).
2. torch.fx— symbolically tracing a model and rewriting its graph programmatically (fusion, quantization, injection).

Topics

• Custom Autograd.Function— forward/backward contracts, ctx, save_for_backward, when to use it.
• Higher-Order Gradients & Jacobians— torch.autograd.grad, create_graph, Hessian-vector products, torch.func.
• fx— Symbolic Model Transformation— tracing, graph rewriting, passes, quant tools.

Key Patterns

• Custom autograd = explicit chain rule, explicit saved tensors.
• create_graph=True unlocks meta-learning / double-backward.
• fx lets you read and rewrite the graph— not just execute it.

Quick Links

• 02 Autograd The Gradient Engine
• 00 Readme