JAX Core Fundamentals: Pure Functional Arrays
Welcome everyone, in this article, I will explore the core philosophy of JAX, it's array system and the critical concept of functional purity.
By the end of this article, you will understand why JAX works the way it does and how to write code that fully leverages its powerful transformations like automatic differentiation, vectorization and compilation.
Step 0: Setting Up the JAX AI stack
Before writing any code, we need to install the modern JAX Stack.
# Install the JAX AI stack (includes JAX, Flax, Optax, and friends)
!pip install -q jax-ai-stack
And in order to verify our installation, we can run the code below
import jax
import flax
import optax
import orbax.checkpoint as ocp
print(f"JAX version: {jax.__version__}")
print(f"Flax version: {flax.__version__}")
print(f"Optax version: {optax.__version__}")
print(f"Orbax version: {ocp.__version__}")
print(f"Available devices: {jax.devices()}")
In the code above, we import and check the different parts of the JAX AI Stack, and the last line of code is to determine what kind of devices we are running on.
Results:
JAX version: 0.8.0
Flax version: 0.12.0
Optax version: 0.2.6
Orbax version: 0.11.26
Available devices: [CpuDevice(id=0)]
The JAX Philosophy: Why purity matters
JAX is designed as a function transformation library. Its transformations like jit, grad, vmap, etc analyze your code like a compiler. For this analysis to hold reliably, your functions must be pure. In the world of programming, these have been known for characteristics such as:-
Pure functions are deterministic i.e. for the same input, the same output will always be gotten
Pure functions are known to have no side effects. i.e. they don't modify state whether local or external scope.
And just like any function, pure funcs rely exclusively on Explicit arguments.
Finally, You should be able to replace a pure function call directly with it's evaluated return value anywhere in the code without changing the program's behavior.
This purity enables aggresive optimizations including XLA compilation for GPUs/TPUs and automatic parallelization.
Writing Pure Functions
In the section below, I'll explore the difference with code blocks
# Impure version
counter = 0
def impure_double(x):
global counter
counter += 1 # Modifies state
return 2 * x
The above codeblock modifies the counter which wasn't passed as a parameter.
# Pure version
def pure_double(x, counter):
new_counter = counter + 1
return 2 * x, new_counter # return everything needed
So the basis being that any data to be modified needs to be passed along. Careful while using condiotnal statements inside pure functions as they should follow the above rules.
Armed with that knowledge base, we can build a reusable linear layer as a pure function, and this is the function for almost all JAX models
def linear_layer(params, x):
"""
params being a PyTree containing weights and bias.
"""
w = params['w'] # shape: (in_features, out_features)
b = params['b'] # shape: (out_features, )
return jnp.dot(x, w) + b
Because this function is pure, we can safely apply transformations to it.
Core JAX Transformations
Automatic Differentiation with jax.grad
Ideally, most of us think that differentiation can be achieved by one of two ways:-
Numerical Differentiation where we approximate gradients using fintie differences(using a mathematical formula), it's easy but prone to floating-point round off errors.
Symbolic Differentiation where we use algebraic manipulation like how we did in high school with characters. It yields the exact formulas but suffers from expression swell leading to complex loops and control flows computationally significant.
Automatic Differentiation (Autodiff) is the middle ground. It breaks down your function into an internal sequence of primitive operations. As it's executed, JAX applies the calculus chain rule to the exact numerical values of these primitives. It computes exact derivatives down to machine precision without generating massive symbolic equations.
In PyTorch or Tensorflow, you compute a forward pass, create a compute a computational graph attached to the data tensors, and call .backward() to mutate a .grad attribute on the variables. JAX rejects this stateful approach. In JAX, differentiation is a pure mathematical function transformation. If you have a function f(x), passing it to jax.grad(f) gives you f'(x) that computes the derivative.
Consider a scalar function below
import jax
import jax.numpy as jnp
def scalar_fn(x):
return jnp.sin(x) ** 2 + 3 * x
# whereas jax.grad is a function transformation, it's a high-order function
grad_fn = jax.grad(scalar_fn)
x_value = 2.0
print(f"Value: {scalar_fn(x_value)}") # Get f(x)
print(f"Gradient: {grad_fn(x_value)}") # Get f'(x)
Code output
Value: 6.826821804046631
Gradient: 2.243197441101074
Consider a more useful use case where we use jax.grad to compute the gradient function.
import jax
import jax.numpy as jnp
# Our main key for randoms
main_key = jax.random.key(0)
def linear_layer(params, x):
return jnp.dot(x, params['w']) + params['b']
def loss_fn(params, x, y_true):
"""Mean square error loss"""
y_pred = linear_layer(params, x)
return jnp.mean((jnp.squeeze(y_pred) - y_true) ** 2) # squeeze to remove singletons
# Create a gradient function wrt the 0th argument (params)
grad_fn = jax.grad(loss_fn, argnums=0)
# Split the main key and keep a copy
w_key, x_key, main_key = jax.random.split(main_key, 3)
# Some data
x = jax.random.normal(x_key, (5, 3)) # Batch 5, 3 features each
y_true = jnp.ones(5)
params = {
'w': jax.random.normal(w_key, (3, 1)),
'b': jnp.zeros(1)
}
# Get the gradient tranformations
gradients = grad_fn(params, x, y_true)
print("Gradient shape for weights: ", gradients['w'].shape)
print("Gradient shape for biases: ", gradients['b'].shape)
Code output:
Gradient shape for weights: (3, 1)
Gradient shape for biases: (1,)
In the above example, I exihibit how JAX seperates model parameters from execution logic. Unlike traditional frameworks, the linear layer maintains no internal state and instead the weights and biases are passed as a nested dictionary which JAX treats as a PyTree. By applying jax.grad with the 0th arg, JAX isolates this parameter tree and traces the execution path to return a grad with a similar structure as the input dict. Finally, jnp.squeeze is to be taken seriously as it prevents mismatched dimensions.
Vectorization with jax.vmap
In standard frameworks, when you compute gradients, the framework automatically sums or averages the gradients across the entire batch to update the model weights. But some ML techniques require the precise gradient for each individual sample in the batch before they are averaged. Pytorch requires workarounds like backward hooks while standard Python for loop forces sequential execution.
Because JAX views differentiation(grad) and vectorization(vmap) as pure transformations, you can nest them inside each other. As opposed to writing your loss function to handle batches, it now calculates the error for a single sample. Then transform it with jax.grad to get a single sample grad function, and finally wrap that function in a jax.vmap to instantly vectorize accross the batch.
Consider the sample Loss function below
import jax
import jax.numpy as jnp
# Obtain our main key for generating randoms
main_key = jax.random.key(0)
def loss_single(params, x, y):
"""Computes MSE loss for a single data point"""
pred = jnp.dot(x, params['w']) + params['b']
return jnp.squeeze((pred - y) ** 2) # Squeeze to ensure it's always scalar
# Transform the gradient function for a single sample
grad_single = jax.grad(loss_single, argnums=0)
# Vectorise the grad_single func over a batch of data
# The in_axes parameters mean:-
# Don't batch the params, but batch x and y along the zero axis
per_sample_grads_fn = jax.vmap(grad_single, in_axes=(None, 0, 0))
# Split the main key and obtain the key for generating the x_batch of sample data
x_batch_key, main_key = jax.random.split(main_key, 2)
params = {
'w': jnp.array([1.5, -2.0, 0.5]),
'b': jnp.array([0.0])
}
x_batch = jax.random.normal(x_batch_key, (4, 3))
y_batch = jnp.array([1.0, -1.0, 0.5, 2.0])
# Now you can calculate gradients for every single sample simultaneously
per_sample_grads = per_sample_grads_fn(params, x_batch, y_batch)
print("Weight gradients shape:", per_sample_grads['w'].shape)
Code output:
Weight gradients shape: (4, 3)
So what's happening with the above code is that with jax.vmap(jax.grad(loss_single)) , JAX executes a two-step transformation before it ever touches the GPU or TPU.
jax.grad traces the single sample function and generates the exact reverse mode auto diff graph required to compute the derivative of one vector.
jax.vmap then intercepts that single sample differentiation graph, then looks at the low-level primitive operations and upgrades them to batch-aware XLA matrix operations.
Instead of running a loop of grads, XLA compiles the entire batch of backward passes into a single parallelized execution block.
Just-In-Time Compilation with jax.jit
In a standard Python script or NumPy code, the Python interpreter executes your script line by line. For an expression like y = jnp.sin(x) ** 2 + 3 * x, Python has to allocate intermediate memory to hold the results per priliminary operation and the round trips are costly.
However, jax.jit intercepts your Python function and passes it to the XLA compiler which looks at the entire function globally and performs operation fusion. This means all operations are registered by the hardware once.
Consider the Loss function below
import jax
import jax.numpy as jnp
import time
# Obtain our main key for generating randoms
main_key = jax.random.key(0)
def loss_fn(x):
"""With our math func from earlier"""
return jnp.sin(x) + jnp.cos(x) ** 2 + 3 * x
# jit transformation from normal Python func
compiled_loss_fn = jax.jit(loss_fn)
# Split key and get data
x_data_key, main_key = jax.random.split(main_key, 2)
x_data = jax.random.normal(x_data_key, (5000, 5000))
# The first call for tracing and compiling
# So first-time exections may be slower
t0 = time.time()
_ = compiled_loss_fn(x_data)
print(f"First call: {(time.time() - t0)*1000:.2f} ms")
# Second call will be machine code
t1 = time.time()
res = compiled_loss_fn(x_data)
jax.block_until_ready(res) # hold exection until the fn execution is done
print(f"Second call: {(time.time() - t1)*1000:.2f} ms")
So jax.jit relies entirely on Tracing. When you invoke a function the first time, JAX swaps out your actual data with an empty placeholder called a Tracer. This knows not your data but rather the datatype and shape dimensions.
Code Results:
First call: 96.79 ms
Second call: 1481.97 ms
Before you step on the table, let me explain. The first timing shows the time taken to compile the function while the second timing shows the time taken to execute the function with our actual data but given the size of our matrix, those are the results, could have been worse.
JAX and Arrays
Now that you're well acquainted with how JAX works, we can easily tackle how it handles arrays.
Immutability
In standard Python and NumPy, arrays are mutable and can be overridden in place. However, in JAX, arrays are absolutely immutable. Once an array is created, it can not be altered and trying will lead to an error and this is because of the JAX design as earlier on explained.
The first solution JAX introduced the .at property helper. Instead of changing the original array, it returns a brand new array with updated values leaving the original array intact.
import jax.numpy as jnp
x = jnp.array([1.0, 2.0, 3.0])
# x[0] = 7.0 -- will fail with a TypeError
# The JAX way
y = x.at[0].set(5.0)
print(f"Original x: {x}")
print(f"New y: {y}")
While creating new arrays may look inefficient, the XLA compiler keeps track of them and if it finds original arrays that were never used again, it optimizes the final machine code with in-place memory mutation.
Duck Typing with NumPy
It should also be noted that JAX implements standard dark typing. This means that jax.numpy mirrors the official numpy API almost perfectly. This also means that a JAX array can be passed directly into many functions exepcting a standard NumPy array. It also exposes the same properties like .shape, .dtype and .ndim
# Creating a NumPy like array
x = jnp.linspace(0, 10, 5)
print(x.shape)
print(x.dytpe)
The Shift
The biggest shift happens in how JAX handles arrays when you wrap your code in a transformation. Remember how JIT replaces actual values with Tracers, out of JIT, a JAX array holds actual values on your GPU/TPU while inside JIT, it's stripped of all that and is merely an abstract tracer. This is also why you can't use raw Python to inspect the elements of a JIT-Compuiled function.
And that's mostly about it for this article.
Notebooks will be available in future articles but I believe basics must be typed out by hand for the fundamental concepts to stick.
Additionally, Wes Kambale (GDE, AI) wrote a beautiful aritlcle about JAX being the NumPy you know but faster. Take a look at his article here now that you know how everyting comes about and experience the power first hand.
See you in the next one
