# Functional State Management and Low-Level Control Flow in JAX

Well be back guys, in our previous article on managing PRNG keys, we explored how JAX abandons the conventional hidden random state in favor of stateless random number generation. We learned that to get deterministic sequences of random numbers, we must pass keys into functions and split them as state updates.

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

```python
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.

![](https://cdn.hashnode.com/uploads/covers/68874bca8a40b70ffef830eb/724096c7-b87a-40b7-8ba9-1409a4c9409b.png align="center")

Consider the Counter example code below

```python
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

```plaintext
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

```python
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.

```python
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

```plaintext
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

```python
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

```plaintext
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

```python
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

```python
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

```plaintext
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

```python
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

```python
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

```plaintext
(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.

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

```plaintext
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.