stringjax_tools.pytrees

Configurable pytree registration helpers for stateful model classes.

JAX reconstructs custom pytree nodes with the registered unflatten function, not by calling __init__. Stateful classes therefore need a deliberate state policy: semantic configuration should be preserved as static auxiliary data, numerical leaves should be traced, and ignored attributes should be limited to recomputable eager caches or scratch state.

The package provides the pytree mechanics only. Consuming packages should define their own static and ignored attribute names locally, usually through a shared PytreePolicy.

State classification

JAX rebuilds custom pytree objects through the registered unflatten function; it does not call __init__. For stateful scientific model classes, classify each attribute before registration:

Attribute kind	Typical treatment	Examples
Static configuration	Preserve with static_keys if hashable and immutable.	model IDs, dimensions, user-supplied scalar bounds
Traced numerical data	Leave as dynamic pytree children.	JAX arrays, differentiable state, model tensors
Recomputable eager cache	Ignore during flattening and restore with ignore_defaults if eager access after reconstruction is expected.	lazy samplers, memoised eigensystems, scratch dictionaries

Do not put semantic state in ignore_keys or ignore_defaults. A common robust pattern is to keep user configuration as static auxiliary data and ignore only the cache built from that configuration. Use ignore_defaults={"_sampler": None, "_cache": dict} when ignored caches must exist after reconstruction; callable defaults are factories for fresh mutable objects.

Static and ignored attribute names must be disjoint. Static values are validated by default: they must be hashable and compare to a scalar, self-equal boolean. This rejects array-like metadata and NaN-like values before they enter JAX auxiliary data.

Pytree policies

PytreePolicy([static_keys, ignore_keys, ...])

Shared pytree flattening policy for one package or model family.

flatten_func(obj, *[, static_keys, ...])	Flatten obj for the JAX pytree protocol.
unflatten_func_class(aux_data, children, ...)	Reconstruct myclass from pytree auxiliary data and children.
make_pytree_flatteners(myclass, *[, ...])	Build class-specific pytree flatten/unflatten callables.