Skip to content

← Flax step 92 of 100

Medium Primitives

flax.struct.dataclass - Pytree-Friendly State

Jump to the task

A training state often groups arrays with values such as a step counter. flax.struct.dataclass gives that state named fields, supports JAX tree operations, and provides a convenient way to create updated instances.

Define a state and replace a field

from flax import struct
import jax.numpy as jnp

@struct.dataclass
class MyState:
    weights: jnp.ndarray
    step: int = 0

old = MyState(weights=jnp.array([1., 2., 3.]))
new = old.replace(step=old.step + 1)

old.step remains zero and new.step is one. Their weights have the same values; replacing the step does not update or copy the weight array. Assigning old.step = 1 raises FrozenInstanceError because fields are frozen by default.

Freezing fields is not deep immutability. If a field contains a mutable Python dictionary or list, that container can still be changed. Use functional updates consistently when state contains nested containers.

What JAX sees

The decorator registers the class as a pytree: a nested structure JAX can flatten into data leaves and reconstruct. By default, fields participate in that traversal. An array field is a leaf; a field containing a dictionary is traversed further. A field is not necessarily a single leaf.

For this MyState, both the weights array and integer step are data leaves. jax.tree_util.tree_map(lambda value: value * 2, new) doubles the weights and step. JAX can treat the integer step as a runtime value under jit.

Dictionaries are also supported pytrees and can be ordinary arguments to jitted functions. They do not need to be hashable for that use. A struct dataclass adds a named schema, methods and field-replacement syntax; it is not required simply to make a dictionary of arrays work with JAX.

Choose static metadata deliberately

A field declared with struct.field(pytree_node=False) is stored as tree metadata rather than traversed as data. For example, a fixed configuration flag used by Python control flow can be static. A boolean can also be a runtime data value when used with suitable array operations or JAX control flow; booleans are not inherently metadata.

Static metadata needs stable equality and hashing for JIT caching. Different metadata values can require different compiled versions, while changing a runtime step value of the same type and shape need not. Keep a frequently changing step counter as data. Do not mark array fields static just because the outer dataclass is frozen: a frozen instance containing arrays is not automatically suitable as a hashable static argument.

The exercise result

Construct a state from params_a, increment the step with replace, and return the weight sum and new step. For [1, 2, 3], the result is [6, 1]. Sum across every dimension of params_a, not only its last axis.

The return value checks the sum and counter, but cannot by itself prove that an implementation used a dataclass or preserved the old instance. When checking your implementation locally, retain both instances and verify that the old step stays zero.

Problem

Implement struct_dataclass_step(seed, params_a, params_b):

  1. Define MyState with @struct.dataclass, weights: jnp.ndarray and step: int = 0. Neither field is static metadata.
  2. Construct it with weights=jnp.asarray(params_a) and step=0.
  3. Use replace(step=instance.step + 1) to obtain the updated state.
  4. Return a JAX array of shape (2,) containing the sum of its weights and its step, in that order. The step is 1.0 in the returned array.

params_a and params_b are finite float32 JAX arrays. Their shapes may differ. The seed is a nonnegative integer-valued scalar that may arrive as a float. Neither seed nor params_b affects the result; they are retained in the function signature. Leave both input arrays unchanged.

Loading visualization…