JAX Programming: From Fundamentals to Large-Scale Systems

A Comprehensive Guide to Accelerated Python Computing

Chapter 1: Why JAX — The Landscape of Python Numerical Computing

1. The Problem with Existing Tools
2. Enter JAX: A New Paradigm
3. What Makes JAX Different
4. Who Should Use JAX (and Who Should Not)
5. How to Read This Book

Chapter 2: Getting Started — Installation, Setup, and Your First Program

1. Python and Environment Prerequisites
2. Installing JAX for Different Hardware
3. Verifying Your Installation
4. Your First JAX Array Operations
5. Understanding the JAX REPL Experience

Chapter 3: JAX Arrays — The Foundation of Everything

1. Creating and Inspecting Arrays
2. Core Array Operations
3. Broadcasting and Reshaping
4. Indexing and Slicing
5. JAX vs NumPy Arrays: Key Differences
6. Memory Layout and Device Placement

Chapter 4: JIT Compilation — Making Python Code Fly

1. What Is JIT Compilation
2. Using @jit in Practice
3. Static vs Dynamic Arguments
4. Tracing and the XLA Backend
5. Common JIT Pitfalls and Solutions
6. Performance Measurement and Profiling

Chapter 5: Automatic Differentiation — The Engine of Modern ML

1. How Automatic Differentiation Works
2. Computing Gradients with jax.grad
3. Higher-Order Derivatives
4. Vector-Jacobian and Jacobian Products
5. Value-and-Gradient Computation
6. Common Autodiff Patterns in ML

Chapter 6: Functional Programming — Embracing JAX Philosophy

1. Why Functional Programming Matters
2. Immutability and No Side Effects
3. Pure Functions and Determinism
4. Closures and Function Composition
5. Thinking Functionally About Data
6. Anti-Patterns to Avoid

Chapter 7: Vectorization with vmap — Writing Loops Without Loops

1. The Batch Processing Problem
2. How vmap Works
3. In-Axis and Out-Axis Specification
4. Nested vmapping
5. Combining vmap with jit and grad
6. When Not to Use vmap

Chapter 8: Parallelism — pmap, pjit, and Distributed Computing

1. Understanding Devices in JAX
2. Device Arrays and Placement
3. Using pmap for Single-Host Parallelism
4. Introduction to pjit
5. Shardings and Mesh Definition
6. Multi-Host Distributed Training

Chapter 9: Control Flow — Loops and Conditionals That Compile

1. Why Python Control Flow Fails Under JIT
2. Conditional Execution with lax.cond
3. While Loops with lax.while_loop
4. For Loops with lax.fori_loop
5. Scan for Sequential Operations
6. Choosing the Right Primitive

Chapter 10: PyTrees — Structuring Complex Data

1. What Is a PyTree
2. Registering Custom Tree Types
3. Flattening and Unflattening
4. Operating on PyTrees
5. PyTrees in Machine Learning Models
6. Advanced PyTree Patterns

Chapter 11: Randomness — Reproducible Stochasticity

1. Why JAX Handles Randomness Differently
2. The PRNG Key System
3. Splitting Keys for Parallelism
4. fold_in for Labeling Keys
5. Seeding and Reproducibility
6. Common Randomness Pitfalls

Chapter 12: Custom Gradients — Beyond Automatic Differentiation

1. When You Need Custom Gradients
2. Using jax.custom_jvp
3. Using jax.custom_vjp
4. Gradient Checking Techniques
5. Practical Examples from Research

Chapter 13: Neural Networks with Flax and Equinox

1. Building Blocks of Neural Networks in JAX
2. Introduction to Flax
3. Introduction to Equinox
4. Flax vs Equinox: Choosing a Library
5. Common Architectures and Patterns

Chapter 14: Optimization — Training at Scale

1. Optimizers in JAX Ecosystem
2. Learning Rate Scheduling
3. Gradient Clipping and Normalization
4. Mixed Precision Training
5. Checkpointing and Experiment Tracking
6. Scaling to Large Models

Chapter 15: Advanced Topics — Scientific Computing, Probabilistic Programming, and RL

1. JAX for Scientific Computing
2. Introduction to Probabilistic Programming with NumPyro
3. Reinforcement Learning with JAX
4. Differentiable Physics Simulations
5. Case Study: Training a Large Language Model

Chapter 16: Performance Engineering and Deployment

1. Profiling JAX Programs
2. Debugging Techniques and Tools
3. Memory Optimization Strategies
4. Interoperability with NumPy and Other Libraries
5. Model Export and Serving
6. Production Deployment Considerations

Conclusion: The Future of JAX and Accelerated Computing

References