Pytree policies

stringjax_tools provides generic pytree mechanics, not package-specific state policy. Each consuming package should define its own static and ignored attribute names, then reuse one policy across the classes that share those conventions.

import jax
import jax.numpy as jnp
import numpy as np

from stringjax_tools.pytrees import PytreePolicy, flatten_func, unflatten_func_class

Shared policy registration

Static attributes travel as JAX auxiliary data. Ignored attributes are dropped from reconstructed traced objects by the base policy, so they must be limited to recomputable caches or scratch state. Semantic state, such as user-supplied bounds or physical parameters, should be static if it is hashable and immutable, or a traced child if it is array-like.

POLICY = PytreePolicy(static_keys=('scale',), ignore_keys=('_cache',))

@POLICY.register
class Model:
    def __init__(self, scale, values):
        self.scale = scale
        self.values = values
        self._cache = {'eager-only': True}

model = Model(scale=2.0, values=jnp.array([1.0, 2.0]))
updated = jax.tree_util.tree_map(lambda x: x + 1.0, model)

print('scale:', updated.scale)
print('values:', updated.values)
print('has cache:', hasattr(updated, '_cache'))

scale: 2.0
values: [2. 3.]
has cache: False

Restoring ignored caches

JAX pytree reconstruction bypasses __init__. If a class method may read an ignored cache after a round trip, use ignore_defaults={...} to restore that cache to a safe value. True configuration remains static, while only the rebuildable cache is ignored and restored.

# The next cell registers a class with restored ignored-cache defaults.
# Callable defaults such as dict would be called once per reconstructed object.

class StatefulModel:
    def __init__(self, h12, user_Q=None, bounds=(-1, 1)):
        self.h12 = h12
        self._user_Q = user_Q
        self._sampler_config = (('bounds', tuple(bounds)),)
        self.values = jnp.ones(h12)
        self._sampler = None

    def Q(self):
        if self._user_Q is not None:
            return self._user_Q
        return self.h12 + 2

    @property
    def sampler(self):
        if self._sampler is None:
            self._sampler = dict(self._sampler_config)
        return self._sampler


STATEFUL_POLICY = PytreePolicy(
    static_keys=('h12', '_user_Q', '_sampler_config'),
    ignore_defaults={'_sampler': None},
)
STATEFUL_POLICY.register(StatefulModel)

stateful = StatefulModel(h12=2, user_Q=100, bounds=(-7, 7))
print('before Q:', stateful.Q())
print('before sampler:', stateful.sampler)

children, treedef = jax.tree_util.tree_flatten(stateful)
rebuilt_stateful = jax.tree_util.tree_unflatten(treedef, children)

print('after has _sampler:', hasattr(rebuilt_stateful, '_sampler'))
print('after _sampler:', rebuilt_stateful._sampler)
print('after Q:', rebuilt_stateful.Q())
print('after sampler:', rebuilt_stateful.sampler)

before Q: 100
before sampler: {'bounds': (-7, 7)}
after has _sampler: True
after _sampler: None
after Q: 100
after sampler: {'bounds': (-7, 7)}

Why configuration should not be ignored

Restoring every ignored attribute to None is not a safe general fix. If an ignored field actually stores user configuration, a round trip can silently change the model. Keep configuration in static auxiliary data and reserve ignored fields for caches that can be rebuilt from the preserved state.

class MisclassifiedModel:
    def __init__(self, h12, user_Q=None):
        self.h12 = h12
        self._user_Q = user_Q
        self.values = jnp.ones(h12)

    def Q(self):
        if self._user_Q is not None:
            return self._user_Q
        return self.h12 + 2


BAD_POLICY = PytreePolicy(
    static_keys=('h12',),
    ignore_defaults={'_user_Q': None},
)
BAD_POLICY.register(MisclassifiedModel)

bad = MisclassifiedModel(h12=2, user_Q=100)
children, treedef = jax.tree_util.tree_flatten(bad)
rebuilt_bad = jax.tree_util.tree_unflatten(treedef, children)

print('before Q:', bad.Q())
print('after Q:', rebuilt_bad.Q())

before Q: 100
after Q: 4

Direct flattening

For lower-level registration flows, the direct flatten/unflatten functions are also available. They take the package policy as explicit arguments, including optional defaults for ignored cache attributes.

children, aux = flatten_func(model, static_keys=('scale',), ignore_keys=('_cache',))
rebuilt = unflatten_func_class(aux, children, Model, ignore_defaults={'_cache': dict})

print('children:', len(children))
print('rebuilt scale:', rebuilt.scale)
print('rebuilt values:', rebuilt.values)
print('rebuilt cache:', rebuilt._cache)

children: 1
rebuilt scale: 2.0
rebuilt values: [1. 2.]
rebuilt cache: {}

Static auxiliary data is validated by default. Non-hashable metadata should not be static unless the consuming package has a clear reason and documents the choice.

class BadStatic:
    def __init__(self):
        self.meta = {'not': 'hashable'}
        self.values = jnp.ones(2)

try:
    flatten_func(BadStatic(), static_keys=('meta',))
except ValueError as exc:
    print(type(exc).__name__ + ':', exc)

ValueError: Static pytree attribute 'meta' is not hashable. Use validate_static=False only if you know JAX can safely carry this value as auxiliary data, or remove it from static_keys.