Overview: The JAX Revolution in Scientific Computing

The Crisis That Changed Everything¶

In Project 4, you spent hours implementing Hamiltonian Monte Carlo. You hand-coded finite difference gradients, watching your computer evaluate the log-posterior thousands of times just to approximate derivatives. It worked. But deep down, you knew: There has to be a better way.

In Project 2, you wrote an N-body simulator. It took minutes to run 100 particles for 100 timesteps. You accepted this as “the cost of doing physics.”

You stand at a crossroads.

Down one path: continue with NumPy, finite differences, CPU-bound loops. The well-trodden path of academic computational physics. It works. Down the other path: automatic differentiation (“autodiff”), just-in-time (JIT) compilation, GPU acceleration. The path that enables previously impractical science. Physics-informed machine learning. Instant parameter inference from surrogate models. Real-time exploration of billion-dimensional parameter spaces.

This module is that crossroads.

By the time you finish Project 5, your N-body code will run 30-100× faster, compute numerically exact gradients automatically (machine precision), and generate 10,000 simulation realizations overnight—training data that will power the machine learning surrogates you’ll build in Module 7.

Learning Outcomes¶

By the end of this overview, you will be able to:

• Explain why JAX exists and how Google’s TPU constraints led to composable transformations

• Identify the three computational eras and where JAX fits in scientific computing evolution

• Connect your Project 4 finite differences pain to automatic differentiation solutions

• Describe the three core JAX transformations (grad, jit, vmap) and what they enable

• Articulate the transformation from script writer to scientific software engineer

Learning Paths & Priorities¶

This overview supports three learning paths. Choose based on your background and time constraints:

What This Module Is About¶

Priority: 🔴 Essential

This module transforms you from script writers to scientific software engineers. You’ll learn to:

• Compute gradients automatically and exactly (no more finite differences!)

• Speed up code by 10-100× through JIT compilation

• Batch thousands of computations effortlessly via vectorization

• Build professional, reusable scientific software (not just scripts that “work”)

• Generate reproducible training data for machine learning (Module 7 foundation)

This isn’t just learning a new library. It’s learning how modern scientific software is built.

By the end of Module 6 + Project 5, you’ll have transformed your N-body simulator from a NumPy script into a professional JAX package that:

• Runs 10-100× faster (seconds instead of minutes)

• Computes forces via automatic differentiation (goodbye, hand-coded derivatives!)

• Batches 10,000 simulations overnight (training data for Module 7)

• Is pip-installable, tested, and documented (portfolio-quality code)

Google’s JAX — Why Does This Library Exist?¶

Priority: 🟡 Important

The Problem Google Faced (2015-2017)¶

Google had built custom AI accelerators called TPUs (Tensor Processing Units) for machine learning. These chips could run computations 10-100× faster than GPUs—but only if software could exploit them properly.

The problem:

• TensorFlow (Google’s first ML framework): Required static computational graphs. You couldn’t write flexible, exploratory Python code—you had to pre-define everything before running.

• PyTorch (Facebook’s framework): Flexible (dynamic graphs) but didn’t compile efficiently to specialized hardware.

• CUDA (NVIDIA’s low-level GPU programming): Full control but required hardware expertise most researchers didn’t have.

What researchers actually wanted: Write ordinary Python code describing physics or ML models, and automatically get:

• ✅ Gradients (for optimization/inference)

• ✅ Compilation (for speed)

• ✅ Hardware acceleration (TPU/GPU support)

• ✅ Parallelization (across multiple devices)

Without needing to be hardware experts or ML engineers.

The Innovation: JAX’s Key Insight¶

JAX’s breakthrough was complete separation of concerns:

You write: Pure Python functions (just describe your computation—your physics)

JAX provides: Independent transformations you compose:

The magic: These transformations compose freely:

# All of these are valid!
jit(grad(f))                  # Fast gradients
vmap(grad(f))                 # Batched gradients
jit(vmap(grad(f)))            # Fast batched gradients
pmap(jit(vmap(grad(f))))      # Multi-GPU fast batched gradients

Why this was novel: Previous frameworks entangled differentiation with execution. You couldn’t easily say “I want gradients of this function, but also compile it and run it on GPU.”

JAX decouples everything — you describe transformations independently, compose them like Lego blocks.

Why It Matters for You (Not Just Google)¶

Google didn’t build JAX for speed alone — they built it because ambitious research questions required automatic differentiation through arbitrary custom code:

• Physics simulators (like your N-body code)

• Domain-specific algorithms (stellar evolution, radiative transfer, hydrodynamics)

• Novel ML architectures (combining data-driven models + physics constraints)

The key insight: If you can differentiate through ANY computation, new science becomes possible.

Current Reality (2025)¶

Frontier research labs (DeepMind, Anthropic, Google Research, national labs) now use JAX as production infrastructure.

You’re learning tools that enable cutting-edge research, not pedagogical simplifications.

JAX wasn’t built for astrophysics — but it solves our problems perfectly.

The Pain Point — Your Project 4 Experience¶

Priority: 🔴 Essential

Remember Project 4?¶

In Project 4, you implemented Hamiltonian Monte Carlo to measure dark energy parameters. For each HMC step, you needed gradients of the log-posterior, which you computed using finite differences — approximating derivatives by evaluating the function at slightly different parameter values.

This worked. You got results. Dark energy parameters estimated. Project submitted.

But the pain was real:

• Scaling nightmare: Each gradient required 2d function evaluations (d = number of parameters)

  • d=2 (your Project 4): 200,000 likelihood evaluations for one HMC run — tolerable

  • d=100 (realistic astronomy problems): 10 million evaluations — completely impractical

• Numerical fragility: Tuning the step size $h$ was finicky

  • Too large: inaccurate derivatives ($\mathcal{O}(h^2)$ error)

  • Too small: floating-point precision errors

  • Optimal $h$ varies by parameter—manual tuning required

• Approximate, not exact: Gradient errors compound through thousands of HMC steps, causing poor mixing and biased estimates

What You’ll Be Able to Do After Module 6¶

Project 4 (Finite Differences):

# Many lines of careful numerical code
grad = np.zeros_like(theta)
for i in range(len(theta)):
    theta_plus, theta_minus = theta.copy(), theta.copy()
    theta_plus[i] += h; theta_minus[i] -= h
    grad[i] = (log_posterior(theta_plus) - log_posterior(theta_minus)) / (2*h)
# Cost: 2d likelihood evaluations (approximate)

After Module 6 (Automatic Differentiation):

grad_log_posterior = jax.grad(log_posterior)
# Cost: ~3-5× forward pass, INDEPENDENT of d (exact!)

One line. Not just shorter — exact (to machine precision), fast (400× faster for d=1000), and automatic (works for any function JAX can trace).

A Taste of JAX — Three Transformations That Change Everything¶

Priority: 🔴 Essential

This is a brief preview of JAX’s core capabilities. Don’t worry about understanding every detail — Parts 1-2 will explain the how and why. For now, see what’s possible.

1. Automatic Differentiation: jax.grad¶

Problem: In Project 2, you hand-coded gravitational forces with nested loops, manual vector math, and careful Newton’s 3rd law bookkeeping.

JAX solution: Write the simpler function (potential energy), get forces for free:

import jax
import jax.numpy as jnp

def gravitational_potential(positions, masses):
    """Compute potential energy U (scalar, simple)."""
    n = len(masses)
    eps =0.01  # Softening length
    G = 6.67e-8  # Gravitational constant in CGS
    U = 0.0
    for i in range(n):
        for j in range(i+1, n):
            r = jnp.linalg.norm(positions[j] - positions[i]) + eps**2.0
            U -= G * masses[i] * masses[j] / r
    return U

# Automatic differentiation: F = -∇U
force_function = jax.grad(gravitational_potential, argnums=0)

What you get: Exact forces via chain rule, automatically. Want to add magnetic fields or modified gravity? Just update U(r) and call jax.grad—forces come free.

Connection to Project 4: This is the same jax.grad you’ll use for HMC gradients—exact, fast, automatic.

2. JIT Compilation: jax.jit — 10-100× Speedups¶

Problem: NumPy loops are slow (Python interpreter overhead). Your Project 2 N-body code took minutes for 100 particles.

JAX solution: Compile Python to optimized machine code:

def compute_forces(positions, masses):
    # ... your N-body force calculation ...
    n = len(masses)
    forces = np.zeros_like(positions)
    for i in range(n):
        for j in range(n):
            if i != j:
                r_vec = positions[j] - positions[i]
                r = np.linalg.norm(r_vec) + 1e-10
                forces[i] += G * masses[i] * masses[j] * r_vec / r**3
    return forces

def compute_forces_jax(positions, masses, G=6.67e-8):
    """N-body forces (JAX version—same logic, jnp instead of np)."""
    n = len(masses)
    forces = jnp.zeros_like(positions)
    for i in range(n):
        for j in range(n):
            if i != j:
                r_vec = positions[j] - positions[i]
                r = jnp.linalg.norm(r_vec) + 1e-10
                forces[i] += G * masses[i] * masses[j] * r_vec / r**3
    return forces

# JIT-compile the JAX version (one line!)
compute_forces_jit = jax.jit(compute_forces_jax)

Now benchmark:

# Setup: 100 particles
n_particles = 100
positions_np = np.random.rand(n_particles, 3) * 1e13  # Random positions [cm]
masses_np = np.random.rand(n_particles) * 1e33        # Random masses [g]

# NumPy timing
start = time.time()
for _ in range(100):
    forces_np = compute_forces_numpy(positions_np, masses_np)
time_numpy = (time.time() - start) / 100

# JAX timing (convert arrays)
positions_jax = jnp.array(positions_np)
masses_jax = jnp.array(masses_np)

# First call: compilation overhead (ignore this one)
_ = compute_forces_jit(positions_jax, masses_jax).block_until_ready()

# Actual timing
start = time.time()
for _ in range(100):
    forces_jax = compute_forces_jit(positions_jax, masses_jax).block_until_ready()
time_jax_jit = (time.time() - start) / 100

print(f"NumPy time: {time_numpy*1000:.2f} ms")
print(f"JAX+JIT time: {time_jax_jit*1000:.2f} ms")
print(f"Speedup: {time_numpy / time_jax_jit:.1f}×")

Typical output (your mileage may vary):

NumPy time: 45.2 ms
JAX+JIT time: 1.2 ms
Speedup: 37.7×

What happened?

JAX compiled your Python function to optimized machine code (via XLA compiler). The XLA compiler:

• Fuses operations (one kernel instead of many Python loops)

• Eliminates intermediate arrays (no memory allocation overhead)

• Optimizes for your specific CPU/GPU (uses SIMD, caching, etc.)

One decorator (@jax.jit) gave you ~40× speedup. This is typical for numerical code.

3. Vectorization: jax.vmap — Batch 1000 Simulations Automatically¶

Problem: You need 10,000 N-body simulations for ML training (Module 1 ensemble averages, computational style). NumPy sequential loop takes hours.

JAX solution: Automatic vectorization over batch dimension:

def simulate(initial_conditions):
    # ... your N-body integration ...
    return final_state

# Automatic batching (no manual loops!)
batched_simulate = jax.vmap(simulate)
results = batched_simulate(initial_conditions_batch)  # 1000 sims in parallel

Impact for Project 5: 10,000 simulations goes from ~3 hours (sequential) to ~3 minutes (GPU batched) = 60× faster. This makes ML training data generation feasible within the semester. Without JAX, you’d generate 1000 sims max—ML models would be data-starved.

The Key: Composable Transformations¶

JAX’s superpower is that these three transformations compose freely:

# Combine them: fast, batched, automatic gradients
fast_batched_gradients = jax.jit(jax.vmap(jax.grad(loss_fn)))

This enables new science—not just faster old science. Workflows that were computationally impossible (10k gradient evaluations for ML training) become routine.

Parts 1-2 explain the how and why. This was just a taste to build excitement.

The Three Computational Eras¶

Priority: 🟡 Important

Scientific computing evolved through three complementary approaches. All three coexist today—you’ll use each depending on the problem:

Module 6 is your bridge from Era 2 to Era 3.

Era 1: Forward Simulation¶

What: Physics equations → simulate forward → observe results

Your course examples: Projects 1-3 (stellar populations, N-body, radiative transfer)

Limitation: Can’t efficiently answer “what parameters produced this observation?” (requires expensive trial-and-error or grid search)

Era 2: Inverse Problems (Bayesian Inference)¶

What: Have observations → infer parameters via probability (Bayes’ theorem)

Your course example: Project 4 (HMC for dark energy parameters from supernova data)

Key tools: MCMC, HMC, nested sampling—all require many gradient evaluations (this is where JAX’s autodiff saves you)

Limitation: Expensive! Each posterior sample requires running the forward model. For complex simulations (hours each), inference takes weeks.

Era 3: Amortized Inference via ML Surrogates¶

What: Train ML models on thousands of simulations → make instant predictions

The breakthrough: Pay upfront cost once (generate training data with JAX’s speed), then inference is cheap forever

Your course arc:

1. Project 5 (Module 6): Generate 10,000 N-body simulations with JAX (jit + vmap makes this feasible)

2. Module 7: Train surrogate model (GP/NN) to learn: parameters → observables

3. Final Project: Use surrogate for instant Bayesian inference (milliseconds instead of hours)

Why JAX enables this: Generating 10k simulations used to take weeks (prohibitive). With JAX’s 10-100× speedups, it’s overnight (feasible). This is the computational bottleneck JAX solves.

Course Arc Through the Eras¶

• Projects 1-3: Era 1 (forward simulation—build physics intuition)

• Project 4, Module 5: Era 2 (inverse problems—learn Bayesian inference) → Pain point: Finite differences for gradients don’t scale

• Module 6 (JAX): Bridge to Era 3 (tools for fast simulation + gradient computation)

• Module 7: Era 3 (train surrogates on simulation ensembles)

• Final Project: Combine all three (physics-informed learning)

Looking Ahead — What You’ll Build¶

Priority: 🔴 Essential

Module 6 Arc¶

Part 1: Conceptual Foundations (🔴) → Functional programming, pure functions, computational graphs

Part 2: Core Transformations (🔴) → grad, jit, vmap (mathematics + practice)

Part 3: N-body Migration (🟡) → Rewrite Project 2 in JAX with autodiff forces

Part 4: JAX Ecosystem (🟢) → Optax, Equinox, Diffrax, NumPyro (enrichment)

Part 5: Professional Software Engineering (🟡) → Packaging, testing, documentation

Part 6: Synthesis (🟡) → Reflection on transformation (script writer → software engineer)

Project 5: Professional JAX N-body Package¶

By the end of Project 5, you’ll have a pip-installable JAX package that:

✅ Runs 10-100× faster (JIT compilation: NumPy 60s → JAX 2s → GPU 0.5s)

✅ Computes forces via autodiff (no hand-coded derivatives):

def U(positions, masses):  # Potential energy (simple!)
    return -sum(G * m_i * m_j / r_ij for all pairs)

forces = -jax.grad(U)(positions, masses)  # Automatic, exact

✅ Batches thousands of simulations (vmap → 10k sims overnight instead of weeks)

✅ Generates ML training data (10k N-body runs → HDF5 dataset for Module 7 surrogates)

✅ Is professionally packaged (tests, docs, CI/CD—portfolio-quality artifact)

The Transformation¶

Before Module 6: Script writer (finite differences, slow loops, one-off NumPy code)

After Module 6: Scientific software engineer (autodiff, vectorization, tested packages)

Paradigm shift:

Approximate methods    →  Exact transformations
One-off scripts        →  Reusable software
Black-box tools        →  Glass-box understanding

This is the bridge to Era 3: Your JAX package enables Module 7 (train GP/NN surrogates) and Final Project (physics-informed ML).

Ready to Begin?¶

Self-check before Part 1:

• Can you explain why JAX exists (Google’s TPU problem)?

• Which computational era do Projects 1-4 belong to?

• How does autodiff solve Project 4’s finite differences pain?

• What do grad, jit, vmap do (conceptually)?

If yes → Continue to Part 1: Conceptual Foundations

Part 1 teaches WHY (functional programming, pure functions, computational graphs). Part 2 teaches HOW (transformation mechanics, composability, debugging).

Continue to Part 1: Conceptual Foundations →