JAXifying Stateful NumPy with Pure Functions

1. The Tale of Two Paradigms: PyTorch vs. JAX

To ground our discussion, let's start with a high-level comparison. PyTorch uses a "tape-based" automatic differentiation system. As your code runs, it dynamically builds a computation graph (the "tape"). When you call .backward(), it traverses this graph in reverse to compute gradients. This approach is flexible and fits naturally with Python's imperative, object-oriented style.

JAX, on the other hand, is a "function transformation" system. It takes an entire Python function, analyzes it, and creates a new, transformed version—a gradient function, a vectorized function, or a just-in-time compiled function.

For JAX's transformations to work reliably and efficiently, the original function must be pure. This is the cornerstone of JAX's design and the key to its performance.

2. What is a Pure Function?

The concept of a pure function comes from functional programming and is analogous to a mathematical function. A function is pure if it satisfies two conditions:

1. Deterministic: Its output depends only on its explicit input values. Given the same inputs, it will always return the same output.
2. No Side Effects: It does not modify any state outside of its own local scope. This means no modifying global variables, no changing input arguments in-place, no printing to the console, and no reading/writing files.

Let's watch a short clip that explains this concept and why it's so important for JAX's compiler.

The key reason JAX enforces purity is because of its tracing-based compilation. When you apply jax.jit, JAX doesn't run your Python code every time. Instead, it runs it once with placeholder "tracer" objects to record the sequence of operations. This sequence is then compiled into highly optimized machine code by XLA. If your function had side effects, they would either be captured incorrectly during that single trace or cause the trace to fail. By guaranteeing purity, you allow JAX to safely analyze and compile your code for maximum performance.

For a deeper dive into practical examples of pure and impure functions, the following reading is excellent.

3. The Solution: Explicit State Threading

So, how do we work with state, which is essential for machine learning (e.g., model parameters, optimizer state), in a purity-driven framework? The answer is the "explicit state threading" pattern.

Instead of a function implicitly modifying an external state (a side effect), we make the state an explicit input and output of the function. The pattern is:

new_state, result = pure_function(current_state, inputs)

The function takes the current state as an argument, computes what it needs to, and returns the updated state along with any other results. The calling code is then responsible for managing this state, typically by replacing the old state with the new one for the next iteration.

The official JAX documentation provides the canonical example of this pattern using a simple counter.

As you saw, the original object-oriented counter, Counter.count(), fails because it modifies self.n as a side effect. The JAX-friendly version, Counter.count(n), becomes a pure function. The responsibility for maintaining the state (n) is moved outside the function.

4. From Counters to Models

This pattern is not just for toy examples; it is fundamental to all JAX programs. Let's apply it to a familiar machine learning task: linear regression. Here, the "state" is the model parameters.

The linear regression example is a microcosm of a standard JAX training loop. The update function is pure. It accepts the current parameters (params), data (x, y), and a learning rate, and it returns a new set of parameters. The training loop itself is responsible for threading this state:

# Pseudo-code for a JAX training loop
params = initialize_parameters() # Initial state

for x, y in data_loader:



  # The update function is pure
  # It takes the current state (params) and returns the new state
  params = update(params, x, y, learning_rate)




# Final state
final_params = params

This explicit passing of state might seem more verbose than the PyTorch equivalent, where an optimizer simply calls optimizer.step() and modifies the model parameters in-place. However, it makes the data flow completely transparent, which simplifies debugging, analysis, and especially parallelization.

5. The Role of Frameworks like Equinox

You might be thinking that manually threading state for all your parameters, optimizer statistics, and layer states (like in batch normalization) could become cumbersome. You are correct, and this is precisely where libraries like Equinox (your chosen framework), Haiku, and Flax come in.

These libraries provide abstractions that handle the state-threading boilerplate for you, allowing you to write code that looks structured and object-oriented, while ensuring it is functionally pure under the hood.

For instance, in Equinox, a training step is often encapsulated in a single pure function, which is then JIT-compiled. Let's look at a concrete example of this pattern from a tutorial on Equinox.

The make_step function is a perfect, practical illustration of today's lesson. It bundles the loss calculation, gradient computation, and parameter update into a single pure function that receives the entire state (model parameters, optimizer state) and returns the new, updated state. This pure function can then be passed to jax.jit for maximum performance.

Key Takeaways

This lesson was a crucial step in shifting your thinking from the imperative PyTorch style to the functional JAX paradigm.

• Purity is Paramount: JAX's function transformations (jit, grad, etc.) require that the functions they operate on are pure.
• Pure Functions: A pure function's output depends only on its inputs, and it has no side effects. This allows JAX to safely trace, compile, and optimize your code.
• Explicit State Threading: To manage state (like model parameters) in a pure way, you must pass the state into your function as an argument and return the updated state as an output. The pattern is (new_state, result) = f(state, *args).
• Frameworks Handle the Boilerplate: Libraries like Equinox are designed to manage this state-threading pattern for you, providing a clean, high-level API while adhering to JAX's functional principles.

You are now equipped with the most important conceptual tool for writing effective JAX code. In our next lesson, we will build on this by exploring the common pitfalls and errors that arise when these principles of purity and immutability are not followed, such as attempting in-place updates or using standard Python control flow inside jitted functions.