Tutorials

User guides are deeper dives into particular topics within JAX that become relevant as your JAX project matures into larger or deployed codebases.

Task API for pure JAX functions

• Logical autosharding for SPMD
• Logical autosharding for MPMD

Task API for Flax modules

• MLP pipeline example with task transforms
• MLP pipeline with external task registration

MultiMesh contexts

• MultiMesh contexts for MPMD