JAX: Composable Transformations for Numerical Computing¶

📚 Table of Contents¶

What is JAX?
Why JAX?
Core Concepts
Learning Path
Chapter Overview
Quick Start

What is JAX?¶

JAX = Autograd + XLA (Accelerated Linear Algebra)

JAX is a Python library that combines:

NumPy-like API - Familiar array operations
Automatic Differentiation - Compute gradients automatically
Just-In-Time Compilation - Transform Python to high-performance machine code
Functional Programming - Composable, predictable transformations
Hardware Acceleration - Run on CPU, GPU, TPU seamlessly

The JAX Philosophy¶

Traditional ML libraries focus on: "What computation to run"
JAX focuses on: "How to transform computations"

JAX provides composable transformations that work together:

import jax
import jax.numpy as jnp

# Simple function
def f(x):
    return x**2 + 2*x + 1

# Automatic differentiation
gradient_f = jax.grad(f)
print(gradient_f(3.0))  # Output: 8.0 (derivative at x=3)

# Vectorization (batch processing)
vectorized_f = jax.vmap(f)
print(vectorized_f(jnp.array([1.0, 2.0, 3.0])))  # Output: [4. 9. 16.]

# Just-In-Time compilation (speed)
compiled_f = jax.jit(f)
print(compiled_f(3.0))  # Compiled and fast

# Compose them all!
fast_grad_batch_f = jax.jit(jax.vmap(jax.grad(f)))

Key Insight: Composable Transformations¶

JAX transforms are first-class functions - you can compose and nest them:

grad(grad(f))           → Compute 2nd derivatives
jit(grad(f))            → Fast gradient computation
vmap(grad(f))           → Gradients for batch of inputs
grad(vmap(f))           → Gradient of vectorized function
jit(vmap(grad(f)))      → All optimizations together

Why JAX?¶

Comparison with NumPy¶

import numpy as np
import jax.numpy as jnp

# NumPy: Great for numerical computing
x_np = np.array([1, 2, 3])
print(x_np ** 2)  # ✅ Works

# But: Can't compute derivatives
# gradient = np.grad(x_np ** 2)  # ❌ Doesn't exist!

NumPy is excellent for numerical computing but lacks automatic differentiation.

Comparison with PyTorch¶

import torch

# PyTorch: Great for ML with auto-grad
x_torch = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
loss = (x_torch ** 2).sum()
loss.backward()
print(x_torch.grad)  # ✅ Gradients work

# But: Stateful (tensors change in-place)
# Harder to reason about (side effects)

PyTorch makes auto-grad easy but uses imperative (stateful) programming.

JAX's Advantages¶

import jax
import jax.numpy as jnp

# JAX: Best of both worlds
x = jnp.array([1.0, 2.0, 3.0])

# ✅ NumPy-like syntax
print(x ** 2)

# ✅ Automatic differentiation (like PyTorch)
grad_fn = jax.grad(lambda x: (x ** 2).sum())
print(grad_fn(x))

# ✅ Pure functions (predictable, composable)
# ✅ JIT compilation (fast, automatic)
# ✅ Device-agnostic (CPU/GPU/TPU)

Why JAX Matters¶

Feature	NumPy	PyTorch	TensorFlow	JAX
NumPy-like API	✅	✅	❌	✅
Auto-Differentiation	❌	✅	✅	✅
Pure Functions	N/A	❌	❌	✅
Composable Transformations	❌	❌	❌	✅
JIT Compilation	❌	Limited	✅	✅
Hardware Agnostic	✅	Limited	Limited	✅
Learning Curve	Easy	Easy	Hard	Medium

Core Concepts¶

1. Functional Programming¶

JAX requires pure functions - functions with no side effects.

# ✅ Pure function
def add(x, y):
    return x + y

# ❌ Impure function (modifies external state)
total = 0
def add_impure(x):
    global total
    total += x
    return total

Why? JAX transformations need predictable behavior.

2. Automatic Differentiation¶

JAX computes derivatives automatically, even complex compositions.

import jax
import jax.numpy as jnp

def f(x):
    return jnp.sin(x) * jnp.exp(-x)

# First derivative
grad_f = jax.grad(f)
print(grad_f(1.0))  # df/dx at x=1

# Second derivative
grad_grad_f = jax.grad(jax.grad(f))
print(grad_grad_f(1.0))  # d²f/dx² at x=1

# Nth derivative
def f_nth_deriv(x, n):
    f_prime = f
    for _ in range(n):
        f_prime = jax.grad(f_prime)
    return f_prime(x)

3. JIT Compilation¶

JAX compiles functions to machine code for speed.

import jax
import jax.numpy as jnp

def slow_function(x):
    result = 0
    for i in range(1000):
        result += jnp.sin(x) * jnp.exp(-x)
    return result

# Without JIT: ~10ms per call
print(slow_function(1.0))

# With JIT: ~1ms per call (after first compile)
fast_function = jax.jit(slow_function)
print(fast_function(1.0))  # First call: compile + run
print(fast_function(1.0))  # Subsequent calls: fast!

4. Vectorization¶

JAX automatically vectorizes functions (batch processing).

import jax
import jax.numpy as jnp

def f(x):
    return x ** 2 + 2*x + 1

# Apply to single input
print(f(3.0))  # 16.0

# Apply to batch automatically
batch_f = jax.vmap(f)
print(batch_f(jnp.array([1.0, 2.0, 3.0])))  # [4. 9. 16.]

# Apply to batch with multiple arrays
def batched_dot(x, y):
    return jnp.dot(x, y)

vectorized_dot = jax.vmap(batched_dot)
x_batch = jnp.array([[1, 2], [3, 4]])  # Shape (2, 2)
y_batch = jnp.array([[5, 6], [7, 8]])  # Shape (2, 2)
print(vectorized_dot(x_batch, y_batch))  # [17, 53]

5. Device Abstraction¶

JAX runs on any device (CPU, GPU, TPU) with same code.

import jax
import jax.numpy as jnp

# Check available devices
print(jax.devices())  # [GpuDevice(id=0), ...]

# Move data to device
x = jnp.array([1, 2, 3])
x_gpu = jax.device_put(x, jax.devices()[0])

# Computation automatically uses device
result = x_gpu ** 2

# Move back to CPU
result_cpu = jax.device_get(result)

Learning Path¶

Level 1: Foundations (2-3 hours)¶

Chapter 1: JAX Basics & Array Operations
Chapter 2: Pure Functions & Functional Programming
Chapter 3: NumPy-like API deep dive

Goal: Comfortable with JAX arrays and functional style

Level 2: Transformations (4-5 hours)¶

Chapter 4: Automatic Differentiation (grad)
Chapter 5: JIT Compilation
Chapter 6: Vectorization (vmap)
Chapter 7: Parallelization (pmap)

Goal: Master individual transformations and compose them

Level 3: Advanced Topics (5-6 hours)¶

Chapter 8: Composing Transformations
Chapter 9: Control Flow in JAX
Chapter 10: Custom Derivatives

Goal: Build sophisticated numerical applications

Level 4: Neural Networks (6-8 hours)¶

Chapter 11: Building Networks with JAX
Chapter 12: Flax Framework
Chapter 13: Training Loops & Optimization

Goal: Train ML models efficiently with JAX

Level 5: Production (4-5 hours)¶

Chapter 14: Performance Optimization
Chapter 15: Device Management
Chapter 16: Debugging JAX Code

Goal: Deploy performant JAX applications

Chapter Overview¶

📖 Part 1: Foundations¶

Chapter 01: JAX Philosophy & Core Concepts

What is JAX, why it exists
NumPy compatibility
Transformations overview
When to use JAX

Chapter 02: NumPy-like API & Arrays

JAX arrays vs NumPy arrays
Operations (element-wise, linear algebra, etc.)
Shape/dtype semantics
Common gotchas

Chapter 03: Pure Functions & Immutability

What are pure functions
Side effects in JAX (don't do this!)
Immutable data structures
Designing functional code

🔄 Part 2: Core Transformations¶

Chapter 04: Automatic Differentiation - jax.grad

How auto-grad works (reverse-mode)
Computing gradients
Jacobians, Hessians, higher derivatives
Multi-argument gradients

Chapter 05: Just-In-Time Compilation - jax.jit

XLA compiler basics
Tracing and compilation
Static vs dynamic shapes
JIT pitfalls and debugging

Chapter 06: Vectorization - jax.vmap

Manual loop elimination
Broadcasting vs vmap
Batching over different axes
Composing with other transformations

Chapter 07: Parallelization - jax.pmap

Multi-device parallelism
Device mesh setup
Collective operations (all-reduce, etc.)
Distributed training basics

🧠 Part 3: Advanced Patterns¶

Chapter 08: Composing Transformations

Nesting transformations
grad(grad(...))
vmap(grad(...))
Common patterns and recipes

Chapter 09: Control Flow in JAX

If statements in JAX
Loops in JAX
scan, map, cond primitives
Conditional gradients

Chapter 10: Custom Derivatives

Implementing custom backward passes
vjp, jvp operations
When to use custom derivatives
Numerical gradient checking

🧠 Part 4: Machine Learning¶

Chapter 11: Building Neural Networks with JAX

Layers and model architecture
Forward/backward passes
Optimizers (SGD, Adam, etc.)
Loss functions

Chapter 12: Flax Framework

Flax as JAX's framework
Module system
Pytrees and TrainState
Common patterns

Chapter 13: Training Loops & Advanced Optimization

Training loop architecture
Gradient clipping, mixed precision
Learning rate schedules
Multi-GPU training

⚡ Part 5: Production & Performance¶

Chapter 14: Performance Optimization

Profiling JAX code
Memory optimization
Reducing compilation time
Microbenchmarking

Chapter 15: Device Management (CPU, GPU, TPU)

Device discovery and setup
Memory management across devices
Distributed arrays
Asynchronous operations

Chapter 16: Debugging JAX Code

Common errors and solutions
Tracing execution
Using print in JIT
Tools and techniques

Quick Start¶

Installation¶

# CPU version
pip install jax jaxlib

# GPU version (NVIDIA CUDA 12)
pip install --upgrade jax[cuda12_cudnn]

# Or with conda
conda install -c conda-forge jax jaxlib

Hello World¶

import jax
import jax.numpy as jnp

# Simple computation
def f(x):
    return x**2 + 2*x + 1

# Evaluate
print(f(3.0))  # 16.0

# Gradient
grad_f = jax.grad(f)
print(grad_f(3.0))  # 8.0

# Vectorization
print(jax.vmap(f)(jnp.array([1.0, 2.0, 3.0])))  # [4. 9. 16.]

# Compiled
fast_f = jax.jit(f)
print(fast_f(3.0))  # 16.0 (but faster!)

# Compose
grad_and_fast = jax.jit(jax.grad(f))
print(grad_and_fast(3.0))  # 8.0 (fast gradient)

Linear Regression Example¶

import jax
import jax.numpy as jnp
from jax import grad

# Generate data
key = jax.random.PRNGKey(0)
X = jax.random.normal(key, (100, 3))
true_w = jnp.array([1.0, 2.0, -1.0])
y = X @ true_w + 0.1 * jax.random.normal(jax.random.PRNGKey(1), (100,))

# Loss function
def loss(w):
    return jnp.mean((X @ w - y)**2)

# Optimize
w = jnp.array([0.0, 0.0, 0.0])
grad_loss = grad(loss)

for i in range(100):
    g = grad_loss(w)
    w = w - 0.01 * g  # Gradient descent

print(f"Learned w: {w}")
print(f"True w: {true_w}")

Neural Network Example¶

import jax
import jax.numpy as jnp
from jax import random, grad

def relu(x):
    return jnp.maximum(0, x)

def predict(params, x):
    for w, b in params:
        x = jnp.dot(x, w) + b
        x = relu(x)
    return x

def loss(params, x, y):
    pred = predict(params, x)
    return jnp.mean((pred - y)**2)

# Initialize random weights
key = random.PRNGKey(0)
params = [
    (random.normal(key, (784, 128)), jnp.zeros(128)),
    (random.normal(random.fold_in(key, 1), (128, 10)), jnp.zeros(10))
]

# Training step
def train_step(params, x, y, learning_rate=0.01):
    grads = grad(loss)(params, x, y)
    return [
        (w - lr * dw, b - lr * db)
        for (w, b), (dw, db) in zip(params, grads)
    ]

# Batch training
X_batch = random.normal(key, (32, 784))
y_batch = random.normal(key, (32, 10))
params = train_step(params, X_batch, y_batch)

Key Takeaways¶

- ┌─────────────────────────────────────────────────────┐
    - JAX at a Glance                        │
  - ┤
│                                                     │
    - What:    NumPy + Autograd + XLA + Functional      │
│                                                     │
    - Why:     Speed + Clarity + Composability          │
│                                                     │
    - How:     Transformations are first-class objects  │
    - - grad()  → Automatic differentiation    │
    - - jit()   → Compile to machine code      │
    - - vmap()  → Vectorize automatically      │
    - - pmap()  → Parallelize automatically    │
│                                                     │
    - When:    Numerical computing + ML + Research      │
│                                                     │
    - Cost:    Requires functional programming style    │
    - Steeper learning curve than PyTorch      │
│                                                     │
  - ┘

Resources¶

Official:

Tutorials:

Community:

Next: Go to 01 Jax Philosophy & Core Concepts to begin learning JAX foundations.

Last Updated: 2026-08-09\ Difficulty: Intermediate (requires Python + NumPy knowledge)\ Estimated Total Time: 25-30 hours