High-Level API for Autograd

on YouTube ↗09m

" summary

New high-level API lives under torch.autograd.fun and exposes concrete functions: jacobian, vector-jacobian product (vjp) and jacobian-vector product (jvp) for first-order, plus hessian and hessian-vector helpers for second-order. Functions take a user function (not its outputs), support requires_grad/create_graph so you can backprop through Jacobians, and pair with a new benchmark suite. Planned work adds real forward-mode AD for jvp/hvp (for node 18), vmap speedups for batch gradients, and better composability with functionalized torch.nn.

AI-written summary · about

" chapters (5)

1. 0:00Introduction to new API
2. 0:46API motivation
3. 3:11API design and features
4. 5:38Usage examples
5. 7:20Future work and feedback

" related talks

• 10mTorchJD: Jacobian Descent in PyTorchPierre Quinton, Valérian Rey · PyTorch Conference Europe 2026
• 11mIntro to JAX: Accelerating Machine Learning researchJake VanderPlas · Google's Machine Learning Virtual Community Day 2021
• 31mHigh performance machine learning with JAXMat Kelcey · PyConline AU 2021
• 39mDifferentiable Functional ProgrammingNoel Welsh · Scala Days New York 2018
• 14mJAX: Accelerated Machine Learning Research via Composable Function Transformations in PythonMatt Johnson · Ray Summit 2021
• 01h 10mJAX: accelerated machine learning research via composable function transformations in PythonMatt Johnson · SPLASH 2020 REBASE
• 19mGetting started with automatic differentiationColin Carroll · PyCon US 2020
• 29mAutomatic differentiation, a tale of two languagesGuillaume Dalle · PyData Global 2024