Resources and Advanced Guides

This section contains examples and tutorials on more advanced topics, such as multi-core computation, automatic differentiation, and custom operations.

Parallel computation

• Distributed arrays and automatic parallelization
• Explicit sharding (a.k.a. “sharding in types”)
• Manual parallelism with shard_map
• Device-local array layout control
• JAX Memories and Host Offloading
• Optimizer State Offloading
• Introduction to multi-controller JAX (aka multi-process/multi-host JAX)
• Fault Tolerant Distributed JAX
• Distributed data loading
• Colocated Python

Machine learning

• The Training Cookbook

Automatic differentiation

• The Autodiff Cookbook
• Custom derivative rules
• Control autodiff’s saved values with jax.checkpoint (aka jax.remat)
• Advanced automatic differentiation

Errors and debugging

• Errors
• Introduction to debugging
• Debugging runtime values
• JAX debugging flags
• Transfer guard

Pytrees

• Custom pytree nodes

Performance optimizations

• Persistent compilation cache
• Buffer donation
• GPU performance tips

Performance benchmarking and profiling

• Benchmarking JAX code
• Profiling computation
• Profiling device memory

Non-functional programming

• Ref: mutable arrays for data plumbing and memory control

External Callbacks

• External callbacks

FFI

• Foreign function interface (FFI)

Modeling workflows

• Gradient checkpointing with jax.checkpoint (jax.remat)
• Ahead-of-time lowering and compilation
• Exporting and serialization

Example applications

• Training a simple neural network, with tensorflow/datasets data loading
• Training a simple neural network, with PyTorch data loading
• Autobatching for Bayesian inference

Deep dives

• Generalized convolutions in JAX
• XLA compiler flags
• JAX Internals: primitives
• JAX internals: The jaxpr language