Skip to main content
Ctrl+K

Advanced automatic differentiation#

Note

This is a placeholder for a section in the new JAX tutorials.

For the time being, you may find some related content in the old documentation:

  • Advanced Automatic Differentiation in JAX

  • The Autodiff Cookbook

  • Control autodiff’s saved values with jax.checkpoint (aka jax.remat)

  • Custom derivative rules for JAX-transformable Python functions.