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