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
• Manual parallelism with shard_map
• Multi-host and multi-process environments
• Distributed data loading

Automatic differentiation

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

Deep dives

• Generalized convolutions in JAX
• List of XLA compiler flags