Functional Updates: From Modify to Return

The biggest mental shift for JAX newcomers is moving from "modify state" to "return new state." In imperative programming, you modify variables in place. In functional programming, you create new values and return them.

The Functional Update Pattern

import jax.numpy as jnp

# IMPERATIVE (NumPy) style — won't work in JAX
# array[i] = value
# dict[key] = value
# object.attribute = value

# FUNCTIONAL (JAX) style — creates new values
# new_array = array.at[i].set(value)
# new_dict = {**dict, key: value}
# new_params = params._replace(attribute=value)  # NamedTuple

.at[] — The Swiss Army Knife

import jax.numpy as jnp

x = jnp.zeros((5, 5))

# Set values at specific positions
x = x.at[0, 0].set(1.0)
x = x.at[2, 3].set(5.0)
x = x.at[4, 4].set(1.0)

# Increment values
counts = jnp.zeros(10, dtype=jnp.int32)
data = jnp.array([3, 1, 4, 1, 5, 9, 2, 6, 5, 3])
# Build a histogram — add 1 at each index
histogram = counts.at[data].add(1)
print(histogram)  # [0 2 1 2 1 2 0 0 0 1]
# Note: duplicate indices accumulate (both 1s add, both 3s add, etc.)

# Scatter operations
indices = jnp.array([0, 2, 4])
values = jnp.array([10.0, 20.0, 30.0])
x = jnp.zeros(6)
x = x.at[indices].set(values)
print(x)  # [10.  0. 20.  0. 30.  0.]

jnp.where: Conditional Updates Without Branching

import jax.numpy as jnp

x = jnp.array([-3.0, -1.0, 0.0, 2.0, 5.0])

# ReLU: replace negatives with 0
relu = jnp.where(x > 0, x, 0.0)
print(relu)  # [0. 0. 0. 2. 5.]

# Clamp between -1 and 1
clamped = jnp.where(x > 1, 1.0, jnp.where(x < -1, -1.0, x))
print(clamped)  # [-1. -1. 0. 1. 1.]

# Conditional based on another array
mask = jnp.array([True, False, True, False, True])
result = jnp.where(mask, x, 0.0)
print(result)  # [-3. 0. 0. 0. 5.]

The Pattern: Carry State Through Returns

import jax
import jax.numpy as jnp

# IMPERATIVE: modify accumulator in a loop
def imperative_cumsum(arr):
    result = []
    total = 0
    for x in arr:
        total += x
        result.append(total)
    return result

# FUNCTIONAL: use jax.lax.scan (functional loop)
def functional_cumsum(arr):
    def step(carry, x):
        new_carry = carry + x
        return new_carry, new_carry  # (next_carry, output)

    _, cumulative = jax.lax.scan(step, 0.0, arr)
    return cumulative

x = jnp.array([1.0, 2.0, 3.0, 4.0, 5.0])
print(functional_cumsum(x))  # [1. 3. 6. 10. 15.]