Getting Started with JAX
Chapter 1: Introduction to JAX
What is JAX?
JAX vs NumPy
Core Design Philosophy: Function Transformations
Installation and Setup
Working with JAX Arrays
Device Management: CPU, GPU, TPU
Practice: Basic Array Operations
Quiz for Chapter 1
Chapter 2: Accelerating Functions with JIT Compilation
The Need for Speed: Why Compile?
Introducing
jax.jitHow JIT Works: Tracing and Compilation
Python Control Flow and
jitStatic vs Traced Values
Common Challenges with
jitHands-on Practical: Applying
jitQuiz for Chapter 2
Chapter 3: Automatic Differentiation with
gradUnderstanding Gradients
Introducing
jax.gradHow Autodiff Works: Reverse Mode
Differentiating with Respect to Arguments
Higher-Order Derivatives (
grad of grad)Value and Gradient (
jax.value_and_grad)Differentiation and Control Flow
Limitations and Considerations
Hands-on Practical: Computing Gradients
Quiz for Chapter 3
Chapter 4: Automatic Vectorization with
vmapThe Concept of Vectorization
Introducing
jax.vmapMapping over Specific Arguments (
in_axes, out_axes)Handling Multiple Batched Arguments
Nesting
vmapCombining
vmap with jit and gradPerformance Considerations for
vmapHands-on Practical: Vectorizing Functions
Quiz for Chapter 4
Chapter 5: Parallelization Across Devices with
pmapIntroduction to Data Parallelism (SPMD)
Introducing
jax.pmapMapping Data to Devices (
in_axes, out_axes)Device Meshes and Axis Names
Collective Operations (
lax.psum, lax.pmean, etc.)Combining
pmap with other TransformationsDebugging
pmapped FunctionsHands-on Practical: Parallel Computation
Quiz for Chapter 5
Chapter 6: Managing State in JAX
Functional Purity and Side Effects
The Challenge of State in Functional Code
Pattern: Explicit State Passing
Using PyTrees for Structured State
Example: Stateful Counter
Example: Simple Optimizer State
Combining State Management with Transformations
Practice: Implementing Stateful Functions
Quiz for Chapter 6
Collective Operations (lax.psum, lax.pmean, etc.)
Was this section helpful?
- Distributed JAX, JAX core contributors, 2024 - Official documentation covering JAX's distributed programming model, including
pmapand collective operations. - JAX: Accelerated Array Programming in Python, Roy Frostig, Matthew James Johnson, Chris Leary, 2018 Journal of Open Source Software, Vol. 3 DOI: 10.21105/joss.00843 - Introduces the JAX library, detailing its design principles for composable function transformations, automatic differentiation, and compilation to XLA for accelerator execution, which underpins parallelization with
pmap. - Horovod: fast and easy distributed deep learning with TensorFlow, Alexander Sergeev, Mike Del Balso, 2018 arXiv preprint arXiv:1802.05799 DOI: 10.48550/arXiv.1802.05799 - Explains the motivation and implementation of efficient collective communication operations, particularly all-reduce, for distributed deep learning training, providing conceptual context for JAX's
lax.psumandlax.pmean. - High Performance Computing: A Practical Approach, Kevin Dowd, Charles Severance, 2018 (O'Reilly Media) - Provides a practical overview of parallel computing concepts, including the SPMD model and the role of collective communication operations, offering a broader context for JAX's approach.