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)
• 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
• XLA compiler flags