Tutorials

• Quickstart
• Key concepts
• Just-in-time compilation
• Automatic vectorization
• Automatic differentiation
• Introduction to debugging
• Pseudorandom numbers
• Working with pytrees
• Introduction to parallel programming
• Stateful computations
• Control flow and logical operators with JIT

Advanced tutorials

• Advanced automatic differentiation
• External callbacks
• Gradient checkpointing with jax.checkpoint (jax.remat)
• JAX Internals: primitives
• JAX internals: The jaxpr language