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
Handling Multiple Batched Arguments
Was this section helpful?
- jax.vmap, JAX developers, 2024 (JAX Documentation) - Provides current, accurate information on
vmap, including usage ofin_axesand examples. - JAX: Composable transformations of NumPy programs, Roy Frostig, Matthew Johnson, and Christopher Leary, 2018 (Google AI Blog) - Introduces JAX's design and how
vmapfits into its system of composable function transformations for numerical computing. - Autograd: Effortless Automatic Differentiation in Python, Dougal Maclaurin, Michael Johnson, and David Duvenaud, 2015 Proceedings of the International Conference on Machine Learning (ICML), Vol. 37 (PMLR (Proceedings of Machine Learning Research)) DOI: 10.5555/3045118.3045239 - Describes Autograd, a predecessor to JAX, outlining automatic differentiation and array programming principles that influence JAX's batching capabilities.
© 2025 ApX Machine LearningEngineered with