Chapter 1: JAX Philosophy & Core Concepts¶

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: automatic differentiation
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: Methods on tensors
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: runs Python interpreter
print(slow_f(3.0))  # Takes ~10ms

# With JIT: compiled machine code
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: manually loop
batch = jnp.array([1.0, 2.0, 3.0])
print(jnp.array([f(x) for x in batch]))  # [2. 5. 10.]

# Batch: with vmap (automatic vectorization)
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: (num_devices, ...)
# Output shape: (num_devices, ...)

(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: returns value based only on input
def add(x, y):
    return x + y

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

# ✅ Pure: no global state
def square(x):
    return x**2

Impure Functions (❌)¶

# ❌ Impure: modifies external variable
counter = 0
def increment():
    global counter
    counter += 1
    return counter

# ❌ Impure: random output (not deterministic)
import random
def random_int():
    return random.randint(0, 100)

# ❌ Impure: I/O side effect
def write_and_return(x):
    with open("output.txt", "w") as f:
        f.write(str(x))
    return x

# ❌ Impure: mutable data (lists modified in-place)
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?¶

Reproducibility: Same inputs always produce same outputs
Parallelization: No race conditions
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 - Learn JAX array operations
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