Skip to content

Chapter 1

Overview

This chapter explores the foundational ideas behind JAX and why it was designed the way it is.

What Makes JAX Different?

Traditional Approaches

Most Python ML libraries follow imperative programming:

# PyTorch (imperative)
import torch

x = torch.tensor([1.0, 2.0, 3.0])
x = x * 2 # Modify in place
x = x + 1 # Another modification
print(x) # [3. 5. 7.]

Problem: Objects change. Tracing what happens is hard.

JAX Approach: Functional Programming

JAX uses functional programming style:

# JAX (functional)
import jax.numpy as jnp

x = jnp.array([1.0, 2.0, 3.0])
y = x * 2 # Returns new array, x unchanged
z = y + 1 # Returns new array, y unchanged
print(z) # [3. 5. 7.]
print(x) # [1. 2. 3.] - still unchanged

Advantage: Predictable, easy to reason about.

Core Principle: Transformations Are First-Class

What's a Transformation?

A transformation takes a function and returns a new function with different behavior:

import jax
import jax.numpy as jnp

# Original function
def f(x):
 return x**2

# Transform it
grad_f = jax.grad(f) # grad() returns a NEW function

# Use the transformed function
print(grad_f(3.0)) # Compute derivative at x=3

Key insight: jax.grad is not a method on f - it's a standalone transformation that creates a new function.

Why This Matters

This design enables function composition:

# Compose transformations
f_transformed = jax.jit(jax.grad(jax.vmap(f)))
# = JIT(Grad(Vmap(f)))
# = (Apply f to batch) → Compute gradients → Compile

Compare with PyTorch's different approach:

# PyTorch
x = torch.tensor(...)
y = x.requires_grad_(True) # Modify tensor
y.backward() # Call method on tensor
print(x.grad) # Access grad attribute

JAX: Transformations are functions, not methods.

-

Key JAX Transformations

1. Automatic Differentiation: jax.grad

Computes gradients automatically (reverse-mode differentiation).

import jax
import jax.numpy as jnp

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

grad_f = jax.grad(f)
print(grad_f(0.0)) # ≈ 1.0 (derivative at x=0)
print(grad_f(jnp.pi)) # ≈ -0.04... (derivative at x=π)

Why automatic? JAX uses the chain rule automatically:

  • sin(x) has derivative cos(x)
  • exp(-x) has derivative -exp(-x)
  • Product rule: (u*v)' = u'*v + u*v'
  • JAX composes all these automatically

2. Just-In-Time Compilation: jax.jit

Compiles Python functions to machine code via XLA.

def slow_f(x):
 result = 0
 for i in range(1000):
 result += x**2 + 2*x + 1
 return result

# Without JIT
print(slow_f(3.0)) # Takes ~10ms

# With JIT
fast_f = jax.jit(slow_f)
print(fast_f(3.0)) # Takes ~1ms (after compilation)

Trade-off: First call slower (compilation time), but subsequent calls much faster.

3. Vectorization: jax.vmap

Automatically vectorizes functions for batch processing.

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

# Single input
print(f(3.0)) # 10.0

# Batch
batch = jnp.array([1.0, 2.0, 3.0])
print(jnp.array([f(x) for x in batch])) # [2. 5. 10.]

# Batch
batch_f = jax.vmap(f)
print(batch_f(batch)) # [2. 5. 10.]

Why better?

  • Clarity: No explicit loops
  • Performance: Vectorized operations are faster
  • Composability: vmap works with transformations

4. Parallelization: jax.pmap

Parallelizes across multiple devices (GPUs/TPUs).

def f(x):
 return x**2

# Parallelize across devices
parallel_f = jax.pmap(f)

# Input shape
# Output shape

(Covered in detail later)

5. Control Flow: jax.cond, jax.lax.scan

Functional versions of if-statements and loops.

# Pure if-statement
def abs_value(x):
 return jax.lax.cond(
 x < 0,
 lambda: -x,
 lambda: x
)

# Pure loop
def sum_n(n):
 def body(carry, x):
 return carry + x, None

 result, _ = jax.lax.scan(body, 0, jnp.arange(n))
 return result

-

Pure Functions: The JAX Requirement

JAX requires pure functions - functions with no side effects.

What's a Pure Function?

A function is pure if:

  1. No side effects: Doesn't modify external state
  2. Deterministic: Same input → same output (always)
  3. No I/O: Doesn't read/write files, print, etc.

Pure Functions ()

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

# Pure
def append_to_list(lst, item):
 return lst + [item] # Returns new list, original unchanged

# Pure
def square(x):
 return x**2

Impure Functions ()

# Impure
counter = 0
def increment():
 global counter
 counter += 1
 return counter

# Impure
import random
def random_int():
 return random.randint(0, 100)

# Impure
def write_and_return(x):
 with open("output.txt", "w") as f:
 f.write(str(x))
 return x

# Impure
def append_to_list(lst, item):
 lst.append(item) # Modifies original!
 return lst

Why JAX Requires Purity

JAX transformations trace your function:

import jax
import jax.numpy as jnp

def f(x):
 return x**2 + jnp.sin(x)

# When you call jax.grad(f):
# 1. JAX traces through f with symbolic values
# 2. Records all operations
# 3. Builds computation graph
# 4. Generates gradient code

grad_f = jax.grad(f)

If your function has side effects, JAX can't trace it properly:

global_counter = 0

def impure_f(x):
 global global_counter
 global_counter += 1 # Side effect!
 return x**2

# This might work the first time, but:
grad_f = jax.grad(impure_f)
# Doesn't track the global_counter modification
# Results are unpredictable

JAX vs NumPy: Immutability

NumPy Arrays (Mutable)

import numpy as np

x = np.array([1, 2, 3])
x[0] = 999 # Modify in place
print(x) # [999 2 3]

JAX Arrays (Immutable)

import jax.numpy as jnp

x = jnp.array([1, 2, 3])
x[0] = 999 # TypeError: JAX arrays are immutable

But you can create new arrays:

x = jnp.array([1, 2, 3])
y = x.at[0].set(999) # Create NEW array with element changed
print(x) # [1 2 3] - original unchanged
print(y) # [999 2 3] - new array

Why Immutability?

  1. Reproducibility: Same inputs always produce same outputs
  2. Parallelization: No race conditions
  3. Tracing: Can build computation graphs safely

NumPy Compatibility

JAX's API is nearly identical to NumPy:

Common Operations

import jax.numpy as jnp
import numpy as np

# Creating arrays
x_jax = jnp.array([1, 2, 3])
x_np = np.array([1, 2, 3])

# Element-wise operations
print(jnp.sin(x_jax)) # [0.84... 0.90... 0.14...]
print(jnp.exp(x_jax)) # [2.71... 7.38... 20.08...]

# Reductions
print(jnp.sum(x_jax)) # 6
print(jnp.mean(x_jax)) # 2.0

# Linear algebra
A = jnp.array([[1, 2], [3, 4]])
b = jnp.array([1, 2])
print(jnp.dot(A, b)) # [5 11]
print(jnp.linalg.inv(A)) # Inverse
print(jnp.linalg.eigvals(A)) # Eigenvalues

Key Differences

# JAX arrays are immutable
x = jnp.array([1, 2, 3])
x[0] = 999 # TypeError!

# Use.at[] indexing
x = x.at[0].set(999)

# Indexing doesn't support fancy indexing like NumPy sometimes
indices = jnp.array([0, 2])
# x[indices] might not work as expected

# Use advanced indexing carefully
x = x.at[indices].set([999, 888])

-

When to Use JAX

JAX is Great For

Research & Experimentation

  • Easy to try new algorithms
  • Composable transformations
  • Clear code

Performance-Critical Code

  • Automatic compilation
  • Multi-device parallelization
  • Memory efficient

Numerical Computing

  • Optimization (gradient descent, etc.)
  • Differential equations (gradients of gradients)
  • Physics simulations

ML Model Development

  • Natural gradient computation
  • Hessian-free optimization
  • Complex architectures

JAX is Not Ideal For

Beginners to ML

  • PyTorch has better tutorials
  • Steeper learning curve
  • More debugging needed

Production Inference Only

  • Overhead for pure research
  • PyTorch/TensorFlow more mature
  • Simpler alternatives exist

Stateful Code

  • Requires functional style
  • No built-in state management
  • (Flax adds this capability)

-

The JAX Programming Model

Flow

Write Python function
 ↓
Apply JAX transformation (grad, jit, vmap, etc.)
 ↓
Trace function with symbolic values
 ↓
Build computation graph
 ↓
Transform the graph (for gradient, compile, vectorize)
 ↓
Generate new function
 ↓
Execute on device (CPU/GPU/TPU)

Example Walkthrough

import jax
import jax.numpy as jnp

def loss(w, x, y):
 """Compute MSE loss"""
 pred = jnp.dot(w, x)
 return jnp.mean((pred - y)**2)

# Create transformed function
grad_loss = jax.grad(loss, argnums=0) # Gradient w.r.t. first arg (w)

# Initialize
w = jnp.array([1.0, 2.0])
x = jnp.array([3.0, 4.0])
y = jnp.array(5.0)

# Call transformed function
gradient = grad_loss(w, x, y)
# JAX computes dL/dw automatically

Summary

Concept Meaning
Transformation Function that takes a function, returns modified function
Pure Function No side effects, deterministic, reproducible
Immutability Data can't be modified after creation
Composition Combining multiple transformations
Tracing Recording operations to build computation graph

-

Next Steps

  • [02 Numpy Like Api & Arrays](/09-jax/01-fundamentals/(02-numpy-like-api-arrays/) - Learn JAX array operations
  • [03 Pure Functions & Immutability](/09-jax/01-fundamentals/(03-pure-functions-immutability/) - Master functional programming in JAX

Checkpoint: Can you explain why JAX requires pure functions? (Hint: tracing)

-

Last Updated: 2026-08-09