Static vs. Dynamic Pytree Fields

1. The Meaning of Static vs. Dynamic

When a JAX transformation like jax.jit encounters a function, it performs a "tracing" step to generate a jaxpr (JAX expression), which is then compiled by XLA. The static/dynamic distinction determines how different parts of your data are handled during this trace.

• Dynamic Fields are the values that flow through your computation. These are typically JAX arrays (or NumPy arrays). During tracing, they are replaced with abstract Tracer objects that record their shape and data type (dtype). The actual numerical values are not known at compile time. The compiled code is therefore generic with respect to the values of these arrays, but specific to their shapes and dtypes. For jax.grad, these are also the (floating-point) arrays with respect to which gradients can be computed.

• Static Fields are treated as compile-time constants. Their concrete values are directly embedded into the jaxpr itself. They can be used to define the structure of the computation, such as controlling Python-level logic (if/else, for loops) that is unrolled during the trace. Static fields are not arrays being transformed; they are part of the function's definition from JAX's perspective.

This has a critical consequence: if you call a jit-compiled function with a different value for a static field, JAX will re-trace and re-compile the function. A change in a dynamic field's values, however, simply re-runs the existing compiled code.

2. The Equinox Convention

Manually managing which arguments are static can be cumbersome, especially with complex models. This is where Equinox introduces an elegant convention that aligns perfectly with its "models are pytrees" philosophy.

Equinox's filtered transformations (filter_jit, filter_grad, etc.) automate the static/dynamic distinction based on a simple rule:

• Dynamic: JAX and NumPy arrays.
• Static: Everything else (Python scalars like int, bool, float, functions, strings, None, etc.).

This convention is applied recursively to all leaves of any pytree passed to the function.

This simple rule is surprisingly powerful. When you build a neural network as an eqx.Module, its weights and biases will be JAX arrays (dynamic), while its activation functions (callable), layer dimensions (int), and dropout rates (float) will be static by default.

What if you have a Python scalar that you need to be dynamic? For instance, a learnable temperature parameter for a softmax function. The solution is to wrap it in a JAX array.

3. Implications for Core Transformations

Let's see how this static/dynamic split plays out across the main JAX transformations, using the Equinox convention.

JIT Compilation (eqx.filter_jit)

Consider a function that includes a Python if statement. For jit to work, the condition must be known at compile time, meaning it must be static.

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

@eqx.filter_jit
def f(x, use_relu: bool):
    if use_relu:
        return jax.nn.relu(x)
    else:
        return -x




# This compiles one version of the function.
f(jnp.array(-1.0), use_relu=True) # Returns array(0.)




# This triggers a RECOMPILATION because the static argument `use_relu` changed.
f(jnp.array(-1.0), use_relu=False) # Returns array(1.0)

In this example, x is a dynamic array, while use_relu is a static boolean. Changing use_relu forces JAX to create a new jaxpr where the if block is resolved differently.

Automatic Differentiation (eqx.filter_grad)

Gradients are computed with respect to floating-point array leaves. Static leaves, or even integer array leaves, are not part of the differentiation process and will have a gradient of None.

When you call eqx.filter_grad on a loss function that takes your model (an eqx.Module pytree) as its first argument, it will automatically traverse the model, find all the floating-point arrays (weights, biases), and compute gradients for them, while correctly ignoring static parts like activation functions or configuration flags.

Vectorization (eqx.filter_vmap)

Vectorization offers another clear example of the split. vmap needs to know which parts of the input to "batch" (add a new axis) and which parts to "broadcast" (reuse across the batch). The Equinox convention provides a perfect default: batch the arrays, broadcast everything else.

In the MLP ensemble example, vmap iterates over a key to create 8 distinct sets of weights and biases (dynamic fields), resulting in a new leading axis on those arrays. However, the activation function within the MLP (a static field) is not replicated; the same function is broadcast and shared across all 8 models in the ensemble. This is both efficient and intuitive.

Conclusion

You have now moved from knowing how to register a pytree to understanding why the static/dynamic split is so important. This distinction is central to JAX's compile-time, functional approach and is what allows it to perform powerful optimizations.

Key Takeaways:

• Dynamic fields are JAX/NumPy arrays whose values flow through a compiled computation. They are the targets of transformations like grad and vmap.
• Static fields are compile-time constants (e.g., functions, booleans, integers) whose values are baked into the structure of the compiled code. A change in a static field triggers recompilation.
• This distinction explains why JAX can unroll loops and resolve if/else branches at compile time, leading to highly optimized code.
• The Equinox library operationalizes this via a simple, powerful convention: arrays are dynamic, everything else is static. Filtered transformations like eqx.filter_jit and eqx.filter_grad handle this partitioning automatically.

In our next lesson, we will put this knowledge to practical use. Now that you understand how a pytree like an Equinox model is structured into dynamic and static parts, you will learn how to perform targeted, functional updates on it using utilities like equinox.tree_at and eqx.filter. This is the key to modifying model state—like applying optimizer updates—within JAX's functional, immutable paradigm.