Getting Started
Reference Documentation
Advanced JAX Tutorials
Notes
custom_vjp
nondiff_argnums
Developer documentation
API documentation
This is a tutorial developed by engineers and researchers at DeepMind.
Tutorials