Manually writing parameter dictionaries ({'w': ..., 'b': ...}) and manually propagating shapes works fine for a two-layer linear regression. But as soon as you build multi-layer perceptrons, convolutional networks or architectures with batch normalization and dropout, handwriting nested dictionary updates becomes tedious and brittle and this is where Flax comes in.

Flax is a high-level neural network library built natively for JAX. It gives you the clean modular ergonomics of object-oriented layer definitions while staying 100% faithful to JAX's functional and stateless philosophy under the hood.

In this post, we will explore how Flax bridges the gap between ergonomic layer design and pure functional execution, how parameter initialization works, and how to use Flax's TrainState abstraction to run a complete training workflow.

Modules as Blueprints, Not State Containers

If you are coming from PyTorch's nn.Module, you are used to a class that holds both the computational architecture and the actual weight tensors inside self as expressed in the code example below

# In PyTorch Architecture and State are tightly coupled in the object
class PyTorchMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 128)  # Weights live inside this attribute
        
    def forward(self, x):
        return torch.relu(self.fc1(x))

In Flax, a linen.Module does not store your parameter tensors inside the instance. Instead, the module is purely a computational blueprint. It defines how inputs should be transformed when parameters are provided.

A Flax module gives you two primary methods:

• module.init(rng_key, dummy_input) which traces the flax module blueprint using a dummy input to allocate and return an immutable dictionary of initialized parameters.

• module.apply(params, input) which takes an explicit parameter dictionary and an input array, runs the forward pass, and returns the result.

Defining Neural Networks with flax.linen

Let's consider a Multi-Layer Perceptron (MLP) with dropout and batch normalization to see how Flax structures layer definitions

import jax
import jax.numpy as jnp
from flax import linen as nn

class ClassifierMLP(nn.Module):
    hidden_dim: int
    num_classes: int
    dropout_rate: float = 0.1

    @nn.compact
    def __call__(self, x:jnp.ndarray, training: bool = False) -> jnp.ndarray:
        """
        Forward pass blue print
        """
        # Layer 1
        # Dense -> BatchNorm -> Relu -> Dropout
        x = nn.Dense(features=self.hidden_dim)(x)
        x = nn.BatchNorm(use_batch_average=not training)(x)
        x = nn.relu(x)
        x = nn.Dropout(rate=self.dropout_rate, deterministic=not training)(x)

        # Layer 2
        # Output Projection
        logits = nn.Dense(features=self.num_classes)(x)
        return logits
        

Before explaining what's going on in the code snippet above, I will draw your attention to the decorator @nn.compact. In standard Python classes, you define the child layers in the __init__ and use them in forward. However, in Flax, the @nn.compact decorator lets you define child layers inside the __call__ method right where they are used.

As for how the layers are working, the nn.Dense(features=self.hidden_dim)(x) linearly projects the input x to the hidden dimension. The nn.BatchNorm(use_running_average=not training)(x) normalizes activations accross the batch. During training when training=True, it calculates the batch statistics and updates running averages. During inference when training=False, it uses the stored running averages. The nnx.relu(x) function applies the standard ReLU activation, max(0, x). The nn.Dropout(rate=self.dropout_rate, deterministic=not training)(x) randomly drops units to prevent overfitting and finally the nn.Dense(features=self.num_classes)(x) linearly maps hidden features to the final class logits.

Initializing and Applying Modules

Because the module is just a blue print, we instantiate the class without passing any tensor data, then initialize its parameters using a PRNG key and dummy input shape.

# Instantiate the blue print
model = ClassifierMLP(hidden_dim=64, num_classes=10, dropout_rate=0.2)

# Create PRNG keys for parameter init and dropout
main_key = jax.random.key(420)
main_key, init_key, dropout_key = jax.random.split(main_key, 3)

# Create dummy batch matching input shape, batch_size=4, input_features=32
dummy_x = jnp.ones((4, 32))

# Initialize parameters and collection variables
variables = model.init(
    {'params': init_key, 'dropout': dropout_key},
    dummy_x,
    training=False
)

# Inspect the returned pytrees
print("Initialized Variable Collections:", variables.keys() )

Output

Initialized Variable Collections: dict_keys(['params', 'batch_stats'])

Notice that the model.init returned a nested dictionary containing two collections i.e. 'params', the learnable weights and biases and 'batch_stats', the running mean and variance tracked by BatchNorm.

To run an inference pass, we pass those variables back to model.apply. Please note that this is a pure and stateless pass.

# Forward inference pass 
logits = model.apply(
    variables,
    dummy_x,
    training=False
)

print("Inference Output Shape: ", logits.shape)

Output

Inference Output Shape:  (4, 10)

Structuring the Training Loop using TrainState

In our previous article, we managed our parameter dictionary and momentum state manually inside a custom carry tuple. In real applications, tracking parameters, optimizer states, batch stats and update rules manually creates boilerplate.

Flax solves this with flax.training.train_state.TrainState.

TrainState is a clean dataclass that bundles your model parameters, optimizer step and gradient update logic into a single immutable JAX PyTree.

Let's integrated Flax with Optax, a standard functional optimization library for JAX to construct a complete training step.

import jax
import jax.numpy as jnp
from flax import linen as nn
from flax.training import train_state
import optax 
from typing import Any

# Define custom TrainState to hold BatchNorm tracking statistics
class CustomTrainState(train_state.TrainState):
    batch_stats: Any

# Setup model, opitimizer and initial state
model = ClassifierMLP(hidden_dim=64, num_classes=10, dropout_rate=0.1)

init_rng = jax.random.key(42)
init_rng, init_param_key, init_drop_key = jax.random.split(init_rng, 3)

dummy_inputs = jnp.ones((8, 32))
init_variables = model.init(
    {'params': init_param_key, 'dropout': init_drop_key},
    dummy_inputs,
    training=False
)

# Setup Adam optimizer using Optax with clipping to avoid exploding gradients 
tx = optax.chain(
    optax.adam(learning_rate=1e-3),
    optax.clip_by_global_norm(1.0) 
)

# Bundle into TrainState
state = CustomTrainState.create(
    apply_fn=model.apply,
    params=init_variables['params'],
    batch_stats=init_variables['batch_stats'],
    tx=tx,
)

The JIT-Compiled Training Step

We now write a pure JIT-compiled train_step. It takes the current state, a data batch, and a PRNG key for dropout, performs the forward/backward pass, updates running statistics, and returns the updated state.

@jax.jit
def train_step(
    state: CustomTrainState,
    batch: tuple[jnp.ndarray, jnp.ndarray],
    dropout_key: jax.Array
) -> tuple[CustomTrainState, jnp.ndarray]:
    """
    Executes a forward + backward + optimizer update step.
    """
    images, labels = batch

    def loss_fn(params):
        # run forward pass passing both params and batch_stats
        # mutable=['batch_stats'] tells Flax to record updated running stats
        logits, mutated_vars = state.apply_fn(
            {'params': params, 'batch_stats': state.batch_stats},
            images,
            training=True,
            rngs={'dropout': dropout_key},
            mutable=['batch_stats']            
        )

        # Cross-entropy loss
        one_hot_labels = jax.nn.one_hot(labels, num_classes=10)
        loss = optax.softmax_cross_entropy(logits=logits, labels=one_hot_labels).mean()

        return loss, (logits, mutated_vars['batch_stats'])

    # Compute loss and gradients with respect to params
    grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
    (loss_val, (logits, new_batch_stats)), grads = grad_fn(state.params)

    # Optax optimizer step that updates params and opt_state inside TrainState
    new_state = state.apply_gradients(
        grads=grads,
        batch_stats=new_batch_stats
    )

    return new_state, loss_val

Running the Training Loop

With train_step compiled, our execution loop remains clean, readable and fast

# Generate a synthetic dataset for demo purposes
dataset_size = 64
batch_size = 8
features = 32

rng_data = jax.random.key(4747)
x_data = jax.random.normal(rng_data, (dataset_size, features))
y_data = jax.random.randint(rng_data, (dataset_size,), minval=0, maxval=10)

step_key = jax.random.key(7)

# Run mini-batch training updates
num_epochs=3
for epoch in range(num_epochs):
    epochs_loss = 0.0
    num_batches = dataset_size

    for i in range(num_batches):
        step_key, subkey = jax.random.split(step_key)
        batch = (
            x_data[i * batch_size : (i + 1) * batch_size],
            y_data[i * batch_size : (i + 1) * batch_size]
        )
        # State in -> State out
        state, loss = train_step(state, batch, subkey)
        epochs_loss += loss
    print(f"Epoch {epoch + 1} | Average Loss: {epoch_loss / num_batches:.4f}")

Outputs

Epoch 1 | Average Loss: 2.7051
Epoch 2 | Average Loss: 2.5244
Epoch 3 | Average Loss: 2.4118

Notice how Flax retains the raw performance and clarity of JAX like how state remains an explicit and immutable object. The train_step is placed and compiled via @jax.jit without hidden side effects. Auxiliary states like batch_stats are tracked and updated deterministically without mutating hidden class properties.

All in all

• Flax linen.Module instances describe computation rather than storing tensor state. Parameters live separately in dictionaries returned by .init().

• Use @nn.compact to define layers and operations directly in the execution path, allowing Flax to automatically manage shape inference.

• Flax's TrainState brings together parameters, optimizer states, step counts, and custom statistics into a single immutable PyTree.

• Flax structures its APIs to ensure your training steps integrate directly with jax.jit, jax.grad, and Optax optimization chains.

Please note that the code above is better explored in a Jupyter Notebook.

See you guys in the next one.

We now zoom out and tackle a broader design pattern that governs all execution in JAX as stated in topic.

For those familiar with Object-oriented ML frameworks like PyTorch or Tensorflow, the mental model of state likely revolves around object attributes that get modified in-place during computation. In JAX, functional purity means that hidden mutations are fundamentally incompatible with compilation with jax.jit and automatic differentiation with jax.grad

In this article, we'll cover how state works in pure functional programming, how to express iterations and conditions using jax.lax primitives and hopefully build a training loop that handles parameters, optimizer state and PRNG keys in a unified state.

Why state must be handled functionally

In traditional code, functions often carry side effects. They may read or mutate global variables, alter internal class properties or perform inplace updates like x+=1. Consider the sample code below

class Counter: 
    def __init__(self):
        self.count=0
    
    def increment(self, amount):
        self.count += amount # This mutates the internal state
        return self.count

As expressed in earlier articles about how JAX transformations like jax.jit and jax.grad work using tracers, if a function relies on side effects or internal mutations, the trace will fail to capture the changes that happen outside the return values or hardcoded values during a single trace.

To avoid all that, JAX requires pure functions that if given the same input, they always produce the same output and above all have zero side effects on state or mutation.

The Explicit State / Carry Pattern

To update parameters, track stats or counters in JAX, we use the Carry Pattern. Instead of updating an object in-place, your function accepts the state as an argument(the carry), performs the required updates using immutable operations and returns the newly computed state alongside any results.

Consider the Counter example code below

import jax
import jax.numpy as jnp

# State management using a pure function
def increment(state: jnp.ndarray, amount: float) -> tuple[jnp.ndarray, jnp.ndarray]
    """This pure function takes the current state, computes the next state and returns both the new state and output """
    new_state = state + amount
    return new_state, new_state

initial_state = jnp.array(0.0)

# Step 1
state, output1 = increment(initial_state, 1.0)
print(f"Output: {output1}, New State: {state}")

# Step 2
new_state, output2 = increment(state, 2.5)
print(f"Output: {output2}, New State: {new_state}")

Output

Output: 1.0, New State: 1.0
Output: 3.5, New State: 3.5

By making state flow through arguments and return values, JAX compilers can optimize operations, parallelize computation across hardware targets and safely differentiate across multi-step algorithms.

Low-Level Control flow primitives with jax.lax

Writing Python conventional flow statements like for loops or if/else statements inside jax.jit compiled functions presents a unique set of challenges. Forexample, a standard for loop with 1000 iterations creates a Jaxpr computational graph with 1000 repeated blocks leading to high compilation overhead. While standard if statements evaluated on tracer values trigger a ConcretizationTypeError because JAX cannot determine truth values of abstract shapes during tracing. To write compiled loops and conditionals that scale efficiently, JAX provides primitives inside the low-level jax.lax module.

Sequential Iteration with jax.lax.scan

jax.lax.scan is JAX's workhorse primitive for sequential and recurrent processing. It executes a loop where an accumulated state, the carry, is passed from step to step while optionally consuming a sequence of inputs and outputting a sequence of results.

Syntax

final_carry, stacked_outputs = jax.lax.scan(
    f,
    init, 
    xs,
    Length=None
)

where

f is the expected step function f(carry, x) -> (next_carry, y) where the carry is the currently accumulated state and x is the single slice of the current iteration step xs[i] and returns a two-element tuple next_carry which is the updated step i+1 and y, the perstep output you may save.

init is the starting carry value before the first iteration begins. This can be a scalar, JAX Array or a PyTree.

xs being the sequence of inputs to iterate through.

The return values are the final_carry which contains the updated state after processing the last element of xs with the same structure and shape as init and stacked_outputs containing all per-step outputs(y) concatenated together along axis 0 with a shape (loop_length, *shape_of_y)

Consider running a cumulative sum

Let's use lax.scan to compute a running sum across an array while tracking the current aggregate state.

import jax 
import jax.numpy as jnp

# step fn
def scan_running_sum(carry: jnp.ndarray, x: jnp.ndarray) -> tuple[jnp.ndarray, jnp.ndarray]:
    """
    carry: running total so far
    x: individual element from input sequence
    """
    new_carry = carry + x
    output = new_carry
    return new_carry, output

xs = jnp.array([1.0, 2.0, 3.0, 4.0, 5.0])
init_state = jnp.array(0.0)

final_carry, running_totals = jax.lax.scan(scan_running_sum, init_state, xs)

print(f"Inputs: {xs}")
print(f"Final carry: {final_carry}") 
print(f"Running Totals: {running_totals}") 

Output

Inputs: [1. 2. 3. 4. 5.]
Final carry: 15.0
Running Totals: [ 1.  3.  6. 10. 15.]

Consider a Recurrent Neural Network (RNN) Step

lax.scan is ideal for sequential models like Recurrent Neural Networks (RNNs) or Transformers processing temporal sequences because it compiles the iteration into a single optimized loop rather than thousands of unrolled nodes. Take an example of the codebase below

import jax
import jax.numpy as jnp

def rnn_cell(hidden_state: jnp.ndarray, x_t: jnp.ndarray) -> tuple[jnp.ndarray, jnp.ndarray]:
  """
  step: h_t = tanh(W_h * h_{t-1} + W_x * x_t)
  fixed identity-like weights are used for demo purposes
  """
  W_h = jnp.array([[0.8, 0.1], [0.0, 0.8]])
  W_x = jnp.array([[1.0, 0.0], [0.0, 1.0]])

  next_hidden = jnp.tanh(jnp.dot(hidden_state, W_h) + jnp.dot(x_t, W_x))
  return next_hidden, next_hidden

sequence_inputs = jnp.array([
    [1.0, 0.5],
    [0.2, 0.8],
    [0.0, 1.0],
    [-0.5, 0.0]
])

init_hidden = jnp.zeros((2,))

# Trace and execute sequence via scan
final_hidden, hidden_states_history = jax.lax.scan(rnn_cell, init_hidden, sequence_inputs)

print(f"Sequence Length: {sequence_inputs.shape[0]}")
print(f"Hidden States Shape over time: {hidden_states_history.shape}")

Output

Sequence Length: 4
Hidden States Shape over time: (4, 2)

Conditional Execution with jax.lax.cond

This is the JIT-compatible equivalent of an if-else statement. This is so because in a typical Python if-else, execution is dynamic based on values evaluated at runtime. Since JAX's jax.jit needs to build a static computational graph, jax.lad.cond allows JAX to trace both branches during compilation and execute only the appropriate branch based on a runtime boolean array.

Syntax

result = jax.lax.cond(
    pred,
    true_fn,
    false_fn,
    operand
)

where

pred is a boolean scalar array like jnp.array(True) or a conditional statement like x>0. It cannot be a multi-element array.

true_fn and false_fn are the signature alignments that must accept a single argument passed via the operand. These are the independent paths based on the evaluation of the pred. If the functions must take multiple arguments, pass them as a tuple. These must also return outputs with identical shapes and dtypes as required by the static compilation graph to keep a fixed shape.

operand is the data to be passed either to the true_fn or false_fn

Consider the gradient clipping example below

import jax
import jax.numpy as jnp

def threshold_activation(x: jnp.ndarray) -> jnp.ndarray:
    """Applies custom piece scaling"""
    pred = jnp.mean(x) > 0.0
    
    true_fn = lambda val: val * 2.0
    false_fn = lambda val: val * 0.1

    return jax.lax.cond(pred, true_fn, false_fn, x)

compiled_activation = jax.jit(threshold_activation)

pos_data = jnp.array([0.5, 1.5, 2.0])
neg_data = jnp.array([-1.0, -0.5, -0.2])

print("Positive batch result:", compiled_activation(pos_data))
print("Negative batch result:", compiled_activation(neg_data))

Output

Positive batch result: [1. 3. 4.]
Negative batch result: [-0.1  -0.05 -0.02]

Dynamic Iterations with jax.lax.while_loop

During compilation, a standard while loop gets unrolled and requires a fixed compile-time iteration count, while a jax.lax.while_loop allows you to run dynamic condition-based loops without unrolling the computation graph.

Syntax

final_val = jax.lax.while_loop(
    cond_func,
    body_func,
    init_val
)

where

cond_func takes the current loop state value starting with the init_val and returns a boolean scalar array i.e. True to keep looping and False to stop.

body_func takes the current loop state value and returns an updated state of the exact same structure, shape and data type.

init_val is the starting state object usually a PyTree passed into a loop.

Consider the code example to find a convergence below

import jax
import jax.numpy as jnp

@jax.jit
def find_next_power_of_two(target):
    # loop state carries 
    # (current_value, step_count)
    init_val = (1, 0)

    # checks if current_value < target
    def cond_func(val):
        current_val, count = val
        return current_val < target

    # doubles current_value and increments count 
    def body_func(val):
        current_val, count = val
        return (current_val * 2, count + 1)

    final_val, total_steps = jax.lax.while_loop(cond_func, body_func, init_val)
    return final_val, total_steps

print(find_next_power_of_two(50)) #(64, 6)

Output

(Array(64, dtype=int32, weak_type=True), Array(6, dtype=int32, weak_type=True))

Note: Because the loop count isn't fixed, reverse-mode automatic differentiation (jax.grad) cannot automatically differentiate through a while_loop without additional bounds.

Combining functional state management with transformations

The power of JAX reveals itself when combining these functional state principles with core transformations i.e jax.jit and jax.grad. Because state updates are represented explicitly as inputs and outputs, wrapping state-modifying functions with JAX transformations is seamless and free of hidden side effects.

import jax
import jax.numpy as jnp
from typing import Dict, Tuple, Any

def compute_loss(params: Dict[str, jnp.ndarray], x: jnp.ndarray, y: jnp.ndarray, key: jax.Array) -> jnp.ndarray:
  """ Compute mean squared error with stochastic noise injection """
  # Inject stochasticity using PRNG key
  noise = jax.random.normal(key, shape=x.shape) * 0.01
  x_noisy = x + noise

  # model forward pass
  predictions = jnp.dot(x_noisy, params['w']) + params['b']
  loss = jnp.mean((predictions - y) ** 2)
  return loss

def train_step(
    state: Tuple[Dict[str, jnp.ndarray], Dict[str, jnp.ndarray], jax.Array],
    batch: Tuple[jnp.ndarray, jnp.ndarray]
    ) -> Tuple[Tuple[Dict[str, jnp.ndarray], Dict[str, jnp.ndarray], jax.Array], jnp.ndarray]:
    """
    Performs a single SGD update with momentum and updates the key.
    
    Carry State: (params, opt_state, key)
    Batch: (x, y)
    """
    params, opt_state, key = state
    x, y = batch

    # Split the keys
    key, subkey = jax.random.split(key)

    # Compute loss and gradient with respect to params (arg 0)
    loss, grads = jax.value_and_grad(compute_loss, argnums=0)(params, x, y, subkey)

    # Momentum SGD Optimizer Logic (Pure functional updates)
    learning_rate = 0.01
    momentum = 0.9  

    new_opt_state = {}
    new_params = {}

    for param_name in params:
      # Update momentum state: v_t = mu * v_{t-1} + g_t
      velocity = momentum * opt_state[param_name] + grads[param_name]
      new_opt_state[param_name] = velocity

      # Update weights: w_t = w_{t-1} - lr * v_t
      new_params[param_name] = params[param_name] - learning_rate * velocity

    new_state = (new_params, new_opt_state, key)
    return new_state, loss


#Initialize inputs, state, and run demo updates
rng = jax.random.key(42)
rng, subkey = jax.random.split(rng)

# Model state initialization
params = {
    'w': jax.random.normal(subkey, (3, 1)),
    'b': jnp.zeros((1,))
}

# Optimizer state
opt_state = {
    'w': jnp.zeros_like(params['w']),
    'b': jnp.zeros_like(params['b'])
}

# Bundle state into carry tuple
state = (params, opt_state, rng)

# Mock Batch Data
x_batch = jnp.ones((8, 3))
y_batch = jnp.ones((8, 1)) * 2.0

# Perform 3 updates
for step in range(3):
  state, loss_val = train_step(state, (x_batch, y_batch))
  print(f"Step {step + 1} | Loss: {loss_val:.4f}")

Output:

Step 1 | Loss: 2.2618
Step 2 | Loss: 1.9221
Step 3 | Loss: 1.3474

If you look closely at the signature and return value of the train step, you'll notice how everything the step needs to know about the current world is packed into a single tuple.

Inside the step, you'll also notice a couple of functional operations happening simultaneously.

• We split the PRNG key so we can safely inject random noise into compute_loss using subkey, while saving key for the next iteration.

• jax.value_and_grad calculates both the scalar loss and the exact parameter gradients with respect to params.

• Instead of modifying weights in-place like param -= lr * grad, we compute brand-new dictionaries for new_opt_state and new_params.

Finally, we pack those updated values back into a new_state tuple and return it alongside the loss scalar.

When we run the outer for loop, we simply overwrite our local state variable with the updated tuple returned by train_step. Because train_step is decorated with @jax.jit, JAX compiles this entire pipeline into a single, highly optimized XLA executable that runs directly on your accelerator with zero side effects.

All in all

1. In JAX, pure functions require state to be explicitly passed in as an argument and explicitly returned as a transformed output (the Carry Pattern).

2. Replace Python for loops in sequential or recurrent workflows with lax.scan to compile long sequences into optimized single-trace instructions.

3. Use JAX-native control flow primitives whenever branching logic or termination conditions depend on dynamic runtime arrays.

4. Combining explicit carry patterns with transformations like jax.jit and jax.grad allows you to write ultra-fast, auto-differentiated training logic without state leakage or compiler errors.

See you guys in the next one.

Those with a background in PyTorch or NumPy, JAX approach to randomness will feel surprisingly strict as there's no global seed or random state. In this article, we shall explore JAX's PRNG system from why global random states break functional purity to understanding key splitting patterns in development. Then later, I'll try to illustrate how to manage them production pipelines and stateful models.

Why JAX rejects global random states

In standard numpy or PyTorch, generating random numbers relies on a global stateful random number generator.

Consider the NumPy code below

import numpy as np

np.random.seed(42)
print(np.random.normal()) # Generates a random number and mutates the global state
print(np.random.normal()) # Generates a different number because a different global state

Output

0.4967141530112327
-0.13826430117118466

Every time you call np.random.normal(), it returns a pseudo-random value and rotates the internal global state counter so the next call produces a different figure. But as established in the first study JAM, JAX functions must remain pure i.e. given the same inputs, the same outputs are expected without any side effects on the output. A hidden mutating global state completely violates the principle of functional purity.

Furthermore, global state makes parallel execution nondeterministic. If two parallel threads request a random number from a single global state, the execution order determines which thread gets which number.

To guarantee reproducibility, parallel safety, and functional purity, JAX makes random states explicit using PRNG Keys.

Abstractions with jax.random.key and jax.random.split

In JAX, randomness is functional. A random function takes a mathematical key array as an explicit argument, deterministically transforms that key, and returns the requested random numbers.

If you pass the exact same key to a random function twice, you will get the same output twice, as demonstrated in the example below

import jax
import jax.numpy as jnp

#create an initial key
key = jax.random.key(42)

#generating randoms with the same key yields the same output
side_a = jax.random.normal(key, (3,))
side_b = jax.random.normal(key, (3,))

print("Draw 1:", side_a)
print("Draw 2:", side_b)
print("Are they identical?", jnp.array_equal(side_a, side_b))

Output

Draw 1: [-0.02830462  0.46713185  0.29570296]
Draw 2: [-0.02830462  0.46713185  0.29570296]
Are they identical? True

Splitting Keys

To generate new, independent random numbers, you must fold or split your existing key into new unique keys using jax.random.split.

import jax

#initialize the root key
root_key = jax.random.key(420)

# split into a usable subkey and hold out the root_key
root_key, subkey = jax.random.split(root_key)

#use the subkey for generating any randoms
new_random = jax.random.normal(subkey, (3,))
print(f"Draw with subkey: {new_random}")

Output

Draw with subkey: [-1.4182527   0.12200668  0.8509691 ]

As a rule of thumb

Never reuse a key. Once a key generates random numbers or splits into subkeys, discard it and use a fresh subkey for subsequent operations.

Managing Keys in Development

When prototyping algorithms or training scripts locally, managing key splitting manually is simple once you establish disciplined coding habits.

Splitting and consuming keys in loops

In sequential scripts or training loops, keep a single key variable in scope. At each step, split the current key into a fresh step_key to consume and update key with the remaining stream.

import jax

root_key = jax.random.key(420)

# training steps with random operations
for step in range(3):
    root_key, step_key = jax.random.split(root_key)

    # pass the key into your model or data transformation
    random_noise = jax.random.normal(step, (2,))
    print(f"Step {step + 1} noise sample: {random_noise}")

Output

Step 1 noise sample: [-1.4182527   0.12200668]
Step 2 noise sample: [-0.06860777  0.04123694]
Step 3 noise sample: [-1.8361075  1.5372876]

Multi-Subkey Splitting

If a single function or initialization step requires multiple random operations, you can split a key into ‭N‬ subkeys simultaneously. Consider the example below

import jax

root_key = jax.random.key(256)

#Generate distinct subkeys at once
subkeys = jax.random.split(root_key, num=4)

print("Subkeys array shape:", subkeys.shape)

Output

Subkeys array shape: (4,)

Vectorizing and Compiling Randomness

Because PRNG keys are standard jax.Array objects under the hood, they fit directly into JAX's core transformations i.e. jit and vmap

Vectorizing Random Sampling with jax.vmap

Suppose you want to generate independent random trajectories or samples across a batch. Instead of running a Python loop, you can split a key into a batch of keys and use jax.vmap

import jax
import jax.numpy as jnp

root_key = jax.random.key(254)

# Get 5 keys
batch_keys = jax.random.split(root_key, num=5)

# Vectorize jax.random.normal across axis 0 of batch_keys
vmapped_normal = jax.vmap(lambda k: jax.random.normal(k, shape=(3, )))

# Get 5 distinct samples of shape (3,) in parallel
batch_samples = vmapped_normal(batch_keys)
print("Batch random samples shape:", batch_samples.shape)

Output

Batch random samples shape: (5, 3)

Randomness inside jax.jit

Passing keys into jax.jit compiled functions works transparently. Since keys are pure arrays, compiling functions that accept keys introduces zero side effects.

import jax
import jax.numpy as jnp

@jax.jit
def apply_dropout(x, key, dropout_rate=0.2):
    keep_prob = 1.0 - dropout_rate

    # Generate binary mask
    mask = jax.random.bernoulli(key, p=keep_prob, shape=x.shape)

    #Scale remaining values to preserve expected value
    return jnp.where(mask, x / keep_prob, 0)

key = jax.random.key(99)
data = jnp.ones((4, 4))

key, subkey = jax.random.split(key)
dropped_data = apply_dropout(data, subkey)
print("Dropped Data:\n", dropped_data)

Output

Dropped Data:
 [[1.25 1.25 1.25 0.  ]
 [1.25 1.25 1.25 1.25]
 [1.25 1.25 1.25 0.  ]
 [1.25 0.   0.   1.25]]

PRNG management when things get serious

While manually splitting with key, subkey = jax.random.split(key) works well for simple scripts, doing this manually across involved neural network architectures becomes unwieldy and error-prone.

Consider the following scenarios from the codebases below

Functional key passing

In pure JAX systems, training steps accept a top-level key as an argument alongside model parameters and batch data. The step function handles the splitting internally and passes dedicated keys down to sub-routines as exhibited in the code sample below.

import jax
import jax.numpy as jnp

def model_forward(params, x, key, is_training=True):
    if not is_training:
        return jnp.dot(x, params['w'])

    w_key, dropout_key = jax.random.split(key)
    out = jnp.dot(x, params['w'])

    # Pass dropout_key to internal dropout logic
    mask = jax.random.bernoulli(dropout_key, p=0.8, shape=out.shape)
    
    return jnp.where(mask, out / 0.8, 0.0)

@jax.jit
def train_step(params, opt_state, x_batch, y_batch, key):
    # Split key for step operations
    step_key, dropout_key = jax.random.split(key)

    # Compute predictions using explicit key
    preds = model_forward(params, x_batch, dropout_key, is_training=True)
    loss = jnp.mean((preds - y_batch) ** 2)

    # ... Gradients and updates ...
    return params, opt_state, loss
 

Stateful PRNG Generators in Class Abstractions

Higher level libraries like Equinox or Flax encapsulate key managementwhen building layer stacks. For example Equinox provides a clean state management pattern but if you're feeling creative, a custom PRNG generator (as exhibited in the previous article) is also an option.

Consider the code sample below

import jax

class PRNGSequence: 
    """ for sequential key splitting """
    def __init__(self, seed_or_key):
        if isinstance(seed_or_key, int):
            self._key = jax.random.key(seed_or_key)
        else: 
            self._key = seed_or_key

    def __call__(self):
        """ returns a fresh subkey while maintaining the internal key state """
        self._key, subkey = jax.random.split(self._key)
        return subkey

# Driver usage
rng = PRNGSequence(seed=420)

# Each call returns a new key
key_0 = rng()
key_1 = rng()
key_2 = rng()

print(f"Key 0 {key_0} - Key 1 {key_1} - Key 2 {key_2}")

Output

Key 0 Array((), dtype=key<fry>) overlaying:
[1499546378 2761259651] - Key 1 Array((), dtype=key<fry>) overlaying:
[3856841681 3513101443] - Key 2 Array((), dtype=key<fry>) overlaying:
[2301777456 4258758523]

By decoupling the key splitting code from model forward passes, your architecture code remains legible while guaranteeing deterministic execution.

All in all

1. JAX has no global seed. All random operations require an explicit PRNG Key array (jax.Array).

2. Calling a random sampler with the same key produces identical numbers. Always use jax.random.split to derive new subkeys.

3. PRNG keys are standard array payloads; they work seamlessly inside jax.jit, jax.vmap, and jax.grad.

4. Use the key, subkey = jax.random.split(key) pattern in loops to keep key streams organized.

5. Pass root keys explicitly into top-level @jax.jit step functions, or wrap key generation in a disciplined abstraction like a PRNGSequence driver or framework modules (Equinox/Flax).

See you all in the next one.

However, in the real world, ML architectures are rarely as simple as single arrays passed into scalar functions. Deep neural networks consist of nested layers, distinct weights and biases, optimizer states and hyperparameter dictionaries.

In this article, we shall move beyond flat arrays and explore how JAX uses Pytrees to represent nested data structures without abandoning functional purity or breaking functional rules. By the end of this article, we shall be able to perform functional parameter updates, stack transformations and leveraging object oriented abstractions.

What are Pytrees and Why do they matter

In JAX, a Pytree is a container-like structure built from standard Python collections that contains leaf elements at its tips

A PyTree typically consists of two main components i.e.

1. Nodes(Branches): These are the structural containers e.g. dicts, lists, tuples, namedtuples or custom registered classes

2. Leaves: This is the actual data payload located at the innermost points of the tree structure. In JAX, these are typically jax.Array objects, NumPy arrays or scalar values.

An example of a PyTree

# Params of a 2-layer network
params = {
    'layer_1': {
        'w': jnp.ones((4, 8)),
        'b': jnp.zeros((8,))
    },
    'layer_2': {
        'w': jnp.ones((8, 1)),
        'b': jnp.zeros((1,))
    },
}

Why PyTrees Matter

In object-oriented frameworks like Pytorch and Tensorflow, model parameters are stored inside stateful class attributes like nn.Module . Gradients are calculated by mutating these internal tensors in place using standard attribute lookup.

But because JAX functions must remain pure, parameters must be seperated from execution logic. PyTrees allow us to bundle arbitrarily complex model parameters into a single structured Python object that can be directly passed into transformed functions.

When you pass a PyTree into functions like jax.grad or jax.jit, JAX automatically traverses the tree structure, performs the required transformations on each individual leaf array and returns a new PyTree with the same container structure as the input.

Basic PyTree Operations

Working with nested structures using standard Python for loops can easily turn verbose and inefficient. JAX provides the jax.tree module and jax.tree_util to interact with PyTree functionality.

Inspecting PyTree Structures and Leaves

Before manipulating a PyTree, it's important to inspect its composition. This can be done by either extracting all leaves into a flat list using jax.tree.leaves or inspecting its container skeleton using jax.tree.structure.

import jax
import jax.numpy as jnp

# A simple param PyTree 
params = {
    'w1': jnp.array([1.0, 2.0], [3.0, 4.0]),
    'b1': jnp.array([0.5, -0.5]),
    'extra': ('learning_rate', 0.01) #Any primitive types are valid tree elements
}

# Extracting Leaves
leaves = jax.tree.leaves(params)
print("Leaves in params: ")
for leaf in leaves:
    print("-", leaf)

# Inspecting the structure
structure = jax.tree.structure(params)
print("\nTree Structure: ")
print(structure)

Output:

Leaves in params: 
- [ 0.5 -0.5]
- learning_rate
- 0.01
- [[1. 2.]
 [3. 4.]]

Tree Structure: 
PyTreeDef({'b1': *, 'extra': (*, *), 'w1': *})

Mapping functions across trees with jax.tree.map

The workhorse of PyTree manipulation is jax.tree.map. It applies a function to every leaf array inside one or more PyTrees, keeping the container architecture untouched.

import jax
import jax.numpy as jnp

params = {
    'w': jnp.array([1.0, 2.0, 3.0]),
    'b': jnp.array([0.5])
}

# Scale all params by a constant factor
scaled_params = jax.tree.map(lambda x: x * 2.0, params)

print("Original params: ", params)
print("Scaled params: ", scaled_params)

Output

Original params:  {'w': Array([1., 2., 3.], dtype=float32), 'b': Array([0.5], dtype=float32)}
Scaled params:  {'b': Array([1.], dtype=float32), 'w': Array([2., 4., 6.], dtype=float32)}

jax.tree.map can also accept multiple PyTrees simultaneously, provided they share the same tree-structural definition. This allows for element-wise operations between matching leaves across distinct trees.

Building a Multi-Layer Neural Network Parameter Tree

Now, let's try to build a multi-layer neural network from scratch using pure functions and PyTrees to illustrate how JAX handles multi-layer architectures.

We shall define parameter initialization logic and forward pass for a Multi-Layer Perceptron(MLP)

import jax
import jax.numpy as jnp

def init_mlp_params(layer_sizes, key):
    params = []
    keys = jax.random.split(key, len(layer_sizes) - 1)

    for in_dim, out_dim, k in zip(layer_sizes[:-1], layer_sizes[1:], keys):
        w_key, b_key = jax.random.split(k)
        scale = jnp.sqrt(2.0 / (in_dim + out_dim))

        weight = jax.random.normal(w_key, (in_dim, out_dim)) * scale
        bias = jnp.zeros((out_dim,))

        params.append({'weights': weight, 'biases': bias})
        
    return params

def mlp_predict(params, x):
    activations = x
    num_layers = len(params)

    for i, layer in enumerate(params):
        outputs = jnp.dot(activations, layer['weights']) + layer['biases']

    # Apply ReLU for hidden layers and linear output for final layer
        activations = jax.nn.relu(outputs) if i < num_layers - 1 else outputs

    return activations

root_key = jax.random.key(42)
# network params:
# 3 input features -> 8 hidden units -> 1 output unit
mlp_params = init_mlp_params([3, 8, 1], root_key)

# Test forward pass with dummy input
dummy_x = jnp.array([1.0, -2.0, 0.5])
prediction = mlp_predict(mlp_params, dummy_x)

print("Initialized network param PyTree")
print(f"Number of layers: {len(mlp_params)}")
print(f"Layer 1 weight shape: {mlp_params[0]['weights'].shape}")
print(f"Prediction for single sample: {prediction}")

Output:

Initialized network param PyTree
Number of layers: 2
Layer 1 weight shape: (3, 8)
Prediction for single sample: [0.4792502]

Combining Transformations on PyTrees

One of JAX's strengths is functional composition, where transformations can be wrapped around each other seamlessly. Suppose we want to compute per gradients accross a batch of data, evaluate loss, and execute the entire operation with compiled XLA performance. We can chain jax.vmap, jax.grad and jax.jit together on functions operating directly on PyTree parameters.

import jax
import jax.numpy as jnp

def compute_loss_single(params, x, y):
    """ MSE for a single sample """
    prediction = mlp_predict(params, x)
    return jnp.squeeze((prediction - y) ** 2)

grad_single_fn = jax.grad(compute_loss_single, argnums=0)

# Vectorize single-sample gradient across batch dimension
per_sample_grads_fn = jax.vmap(grad_single_fn, in_axes=(None, 0, 0))

compiled_per_sample_grads = jax.jit(per_sample_grads_fn)

# Generate batched data
root_key = jax.random.key(101)
x_batch = jax.random.normal(root_key, (5,3))
y_batch = jnp.array([1.0, 0.0, 1.0, 0.5, -1.0])

per_sample_grads = compiled_per_sample_grads(mlp_params, x_batch, y_batch)

print("Per-sample weights gradient shape for Layer 1:", per_sample_grads[0]['weights'].shape)
print("Per-sample biases gradient shape for Layer 1: ", per_sample_grads[0]['biases'].shape)

Output

Per-sample weights gradient shape for Layer 1: (5, 3, 8)
Per-sample biases gradient shape for Layer 1:  (5, 8)

By nesting jax.vmap around jax.grad and compiling the pipeline with jax.jit, JAX avoids the standard Python overhead and converts the entire execution into a single GPU/TPU accelerated compute kernel.

Functional Parameter Updates with jax.tree.map

Since JAX arrays are strictly immutable, we cannot perform inplace updates like params -= learning_rate * grad. Instead, model optimization is expressed as a pure function that inputs current parameters and compute gradients returning an updated parameter PyTree.

We achieved this cleanly across arbitrarily nested parameter structures using jax.tree.map

import jax
import jax.numpy as jnp

def compute_batch_loss(params, x_batch, y_batch):
    """MSE averaged over a batch """
    predictions = jax.vmap(mlp_predict, in_axes=(None, 0))(params, x_batch)
    return jnp.mean((jnp.squeeze(predictions) - y_batch) ** 2)

@jax.jit
def train_step(params, x_batch, y_batch, learning_rate=0.05):
    """Compute batch loss, calculates the batch gradients and applies SGD updates"""
    loss, grads = jax.value_and_grad(compute_batch_loss)(params, x_batch, y_batch)

    # Perform functional gradient descent update across all PyTree leaves
    updated_params = jax.tree.map(
        lambda p, g: p - learning_rate * g,
        params,
        grads
    )

    return updated_params, loss


# Our training week
root_key = jax.random.key(420)
x_train = jax.random.normal(root_key, (100, 3))
y_train = jnp.sin(x_train[:, 0]) + x_train[:, 1]

current_params = init_mlp_params([3, 16, 1], root_key)
print("Training MLP Model using Functional PyTree Updates...")
for step in range(5)
    current_params, loss_val = train_step(current_params, x_train, y_train)
    print(f"Step {step + 1} | Batch Loss: {loss_val:.4f}")

Output:

Training MLP Model using Functional PyTree Updates...
Step 1 | Batch Loss: 1.0306
Step 2 | Batch Loss: 0.7979
Step 3 | Batch Loss: 0.6517
Step 4 | Batch Loss: 0.5486
Step 5 | Batch Loss: 0.4708

Notice how jax.tree.map seamlessly traverses all weight and bias arrays across layers in a single line, subtracting the scaled gradient without requiring explicit manual key lookups or conditional checks.

Introduction to Equinox

While managing parameters manually via Python dictionaries and lists is educational, large production neural networks with complex layers can become tedious to organize using raw nested dictionaries.

To handle this, the JAX library offers libraries like Equinox. Equinox introduces class-based module definitions without breaking the purity rules.

In Equinox, model modules are Python dataclasses registered as PyTrees under the hood. Array attributes are treated as PyTree leaves, while non-array attributes like metadata or configuration integers are treated as static structure nodes.

import jax
import jax.numpy as jnp
import equinox as eqx

# Class based module
class EquinoxMLP(eqx.Module):
    layers: list
    
    def __init__(self, key):
        keys = jax.random.split(key, 2)

        self.layers = [
            eqx.nn.Linear(in_features=3, out_features=16, key=keys[0]),
            eqx.nn.Linear(in_features=16, out_features=1, key=keys[1])
        ]

    
    def __call__(self, x):
        x = jax.nn.relu(self.layers[0](x))
        return self.layers[1](x)

# Instantiate Equinox model module
key = jax.random.key(7)
model = EquinoxMLP(key)

# The model instance is a valid PyTree
leaves = jax.tree.leaves(model)
print(f"Total parameter leaf arrays in model: {len(leaves)}")

# Test forward evaluation
sample_input = jnp.array([0.5, 1.2, -0.8])
output = model(sample_input)
print(f"Model output: {output}")

Output

Total parameter leaf arrays in model: 4
Model output: [-0.37182418]

Because an eqx.Module instance is a PyTree, you can pass the model object directly as the parameter argument into jax.grad, jax.jit, and jax.vmap without extra wrapper dictionaries or boilerplate logic.

In Summary

1. A PyTree is any nested structure (dicts, lists, tuples) whose leaves are array data. They allow JAX to process complex parameters in pure functional transformations.

2. You can apply function operations to every leaf in a nested parameter architecture simultaneously while maintaining tree definitions.

3.Advanced combinations like vmap(grad(fn)) work seamlessly on PyTree structures, allowing per-sample gradient computation and fast parallel processing.

4. Model optimization is performed by passing current parameter trees and gradient trees into jax.tree.map, producing updated parameter trees without in-place mutation.

5. Higher-level frameworks like Equinox rely on PyTrees under the hood, enabling clean, class-based neural network code that integrates directly with native JAX transformations.

See you guys in the next one.

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:-

1. Numerical Differentiation where we approximate gradients using fintie differences(using a mathematical formula), it's easy but prone to floating-point round off errors.

2. 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

At it’s core, the JAX AI Stack extends JAX, the numerical computing library with a suite of loosely coupled Google-backed libraries. Rather than a single monolithic framework, it provides best-in-class tools for each stage of the ML Lifecycle. This design has profound advantages such as:-

• Iterative evolution without breakage:- Data pipelines, checkpointing, and optimization can be updated independently without destabilizing the core numerics engine.

• Composability:- Developers mix and match components to fit specific needs, from rapid experimentation to hero-scale training.

• Durability:- The core JAX library remains focused and adaptable for future hardware and alogorithmic shifts.

This philososphy makes the stack resilient in a field where architectures converge but optmimzation techniques continue to advace quite fast.

Core Compents

The JAX AI Stack covers tge full production pipeline with battle-tested libraries:-

1. JAX:- This is the foundation. It offers a NumPy-like API with powerdul transformations like jit for compilation, grad for differentiation, vmap for vectorization and shard_map for parallelism. Its compiler-first design leverages XLA for aggressive whole-prorgam optimizations and seamless caling accross TPU pods

2. Flax:- This is a flexible neural network authoring. It provides an intuitive, object-oriented(NNX) API on top of JAX’s functional core, making it easy to build, debuf, modify and combine models which is crucial for techniques like LoRA and quantization.

3. Optax:- This is used for composable optimization. It delivers modular gradient transformations, losses and optimzers that chain together declaratively. This enables complex strategies with minimal code while containing scalability correctness.

4. Grain:- This is for deterministic, scalabale data pipelines that are checkpointable and reproducible which is essential for experiments at scale.

Additional infrastracture like the XLA compiler and Pathways for orchestrating computation accross alot of multiple chips underpins extreme scale.

Performance at the Edge

As models mature, peak performance increasingly relies on megakernels which are highly optimized low-level implementations that maximize hardware utilization through overlapping compute, memory and communication. The JAX ecosystem addresses this head-on with a continuum of abstraction levels:-

• High-level XLA optimizations for automated gains.

• Pallas for writing custom kernels directly in Python.

• Tokamax for state-of-the-art kernels

• Qwix for non-intrusive quantization to boost speed and efficiency

• XProf for deep, hardware-aware profiling

This range lets teams with high productivity drill down to expert-level control without leaving the ecosystem, as foundation models demand every last ounce of efficiency on accelerations like TPUs

From Research to Production

The stack isn’t just for training; it supports the entire journey:-

• MaxText / MaxDiffusion:- Scalable reference implementations for LLMs and diffusion models, demostrating production-grade patterns.

• Tunix:- Advanced post-training and alignement

• Inference:- Tight integration with cLLM on TPUs and a dedicated JAX serving runtime for high-throughput, low-latency deployment

Real-world scale is already proven, with massive distributed training runs on tens of thousands of TPUs powered by this stack.

Why It's Built for the Future

The JAX AI stack excels for foundation models because it aligns with their core requirements.

1. Scale works:- It has native support for massive parallelism and orchestration across multiple TPU clusters.

2. Designed for modern hardware:- it has a deep integration with TPUs via XLA and Pathways, delivering specialized performance while remaining open-source and portable.

3. Accelerates research and is reliable at scale:- It's designed based on functional programming, composability, and modular libraries that accelerate iteration while providing tested, scalable primitives.

4. It's here to stay:- it's absraction continuum and loose coupling prepare it for megakernel trends, new model architectures and evolving hardware without having to write the entire codebase

5. Built in the Open:- Everything is Open-source, encouranging contributions and transparency.

In the current era where foundation models drive breakthroughs but also demand unprecedented resources, the JAX AI stack strikes an elegant balance between high and low-level power. Whether you're training the next Gemini-scale model or deloying inference, it provides a robust, forward-looking platform on Google's Cloud TPUs.

For developers and organizations ready to tackle the next wave of AI, exploring the JAX AI Stack is a startegic move towards building systems that are not only powerful today but adoptable for whatever comes next.

git reset --hard origin/master
rm -rf node_modules # The JS guys need this for good luck

Then go touch grass. Because what you’re about to do is the software industry’s version of yelling “Hold my beer” right before natural selection happens.

The Sacred Law

There’s an unwritten law in software engineering

Thou shall not deploy on Friday after 2p.m unless thou hatest the entire team on-call

This law is older than Docker, older than microservers, even older than companies that have their root branch as “main”. It was written in stone right after “Do not rewrite the entire backend in Rust the week before Black Friday.”

The Five Stages of a Friday Deployment

1. Optimism (2:13p.m)

   “It’s just a one-line change. Literally one line. The tests passed. I can even merge this in my sleep.”

2. Hubris (2:27 pm)

   “This is small, let me make the change ASAP. The evening meeting will swing in my favour. My portrait will hang in the office kitchen, right next to the water dispenser.” or so you tell yourself while patting yourself on the back.

3. The Incident (3:02 p.m)

   Production is now returning 502s. The database is speaking in tongues. The monitoring dashboard looks like a Christmas tree that swallowed a rave. Slack is on fire. The engineering manager is typing “what did you do” with increasing aggression.

4. Bargaining(4.45 p.m)

   You’re about to start the usual “But it’s working fine on my locallhost”. You’re promising the world you’ll now read the “Terms and Conditions” before clicking ‘Agree’ if the logs would give you a clue. You consider sacrificing an intern to the DevOps gods.

5. Weekend through the roof (6:00 p.m - Sunday 11:59 pm)

   You are now legally married to PagerDuty. Your significant other has left you for someone who has weekends. Your cat misses you. Your plants are dead. You haven’t seen sunlight since the release of that first AI-generated video of Will Smith eating spaghetti.

Quotes from Developers who deployed on Friday

• “It was literally a typo in the config” - Famous last words, June 2022

• “The CI passed, why wouldn’t production be fine?” - Man who now sleeps in office, since mid 2025

Rational excuses to avoid Friday deploys

• Sudden religious conversion (you now observe “No-Deploy Jumu’ah”)

• Your cat is having an existential crisis and needs emotional support

• Mercury is in retrograde and Jenkins is moody lately

My thoughts

1. If it can wait until Monday, it must wait until Monday.

2. “it’s just a one-liner” is how every war crime in software history started

3. Schedule the deploy for Monday 9:05a.m. Then immediately book a “dentist appointment” for 9:00-11:00a.m. You’re welcome.

Deploying on Friday isn’t brave. It isn’t dedication to the craft. It’s just performance art titled “How to convert weekend plans into PTSD”. And if you absolutely must deploy on Friday, atleast update your LinkedIn to “Open to Work” first.

Thank you very much for your attention to this matter. Now go forth and sin no more…

I have been a backend guy for some time now, dealing with APIs, queues, databases, and the usual grind that keeps systems running. Now I’m on a deliberate mission to level up my entire stack starting with foundation technologies that many people skip, then build all the way up to production-grade machine learning.

This blog is my live journal of that journey.

No polished courses, no fake expert vibes, just me figuring it out step by step, sharing some notes, dumb mistakes, and the “Hmm, that actually worked” moments when they happen.

What You’ll See Here

I am tackling this in a very intentional order:

1. Cloud Networking:

I realised most cloud problems like latency, costs, outages, security etc, trace back to networking. So I’m going hard on VPCs, peering, load balancers, firewalls, private service access, hybrid connectivity, and all the invisible plumbing that makes or breaks production cloud environments. Expect lots of “why is this taking so long?” Posts, terraform configs that finally fixed my egress bills, and stories from when I accidentally open-sourced my entire subnet (no, it won’t happen😂).

2. Cloud Computing:

Once networking feels rock-solid, I’ll move into the broader cloud picture. I’ll take on the real horror stories of IAM, observability, cost optimisation, serverless vs GKE vs VMs, disaster recovery, and how to actually design systems that don’t explode at 3 a.m. This is where I’ll start connecting the dots between backend habits and cloud-native thinking.

3. AI/ML and MLOps:

After I’m confident in the infrastructure, I will go all-in on machine learning and MLOps. I want to build production-ready pipelines on Vertex AI, Kubeflow, feature stores, model monitoring, and everything that makes AI actually shippable and not just Jupyter notebooks. By the time I get there, I’ll have the networking and cloud foundation to understand why things break in production, not just how to train a model.

What to expect along the way

• Spontaneous updates: I’ll post when I learn something, not when the calendar says so.

• First-person battle stories: If I run into something and I find it interesting after solving it, I am definitely sharing that one.

• Code, configs, and commands: Copy-paste ready snippets with explanations of what I tried and failed before.

• Zero fluff: I have no intention of selling courses, and I’m not pretending to have all the answers. I am just documenting my past so you can skip the same pain.

Why I’m sharing this publicly

• Writing forces me to understand the material.

• If my mistakes and breakthroughs can save even one person the same late-night Google rabbit holes, it’s worth it

• I want to look back in a year and see how far I’ve come

How to follow the Journey

• New posts will drop whenever something worth sharing happens (no fixed schedule).

• I’ll cross-post or share summaries on LinkedIn, X, dev.to, or Reddit

• If you want to ride along, subscribe here, on X or LinkedIn so you can know as soon as I drop something

Now that this is known, let me get cooking. See you in the next one