JAX¶

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
x_np = np.array([1, 2, 3])
print(x_np ** 2) # Works

# But
# 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
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
# 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
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
print(slow_function(1.0))

# With JIT
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](/09-jax/01-fundamentals/(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