stringjax_tools.pytrees.PytreePolicy

class PytreePolicy(static_keys=(), ignore_keys=(), ignore_defaults=None, static_types=(<class 'str'>, <class 'bool'>), validate_static=True)

Bases: object

Shared pytree flattening policy for one package or model family.

A StringJAX package should define its domain-specific static_keys and ignored-cache policy locally, then reuse one policy to register all classes that share the same state conventions.

Warning

Ignored attributes are suitable for recomputable caches, scratch arrays and eager-only helpers, but not for semantic state such as user-supplied bounds, physical parameters or configuration values. Put such state in static_keys when it is hashable and immutable, or keep it as a traced child when it is array-like. If an ignored cache may be read after reconstruction, declare an ignore_defaults entry so unflatten restores a safe value such as None or a fresh dict.

__init__(static_keys=(), ignore_keys=(), ignore_defaults=None, static_types=(<class 'str'>, <class 'bool'>), validate_static=True)

Methods

__init__([static_keys, ignore_keys, ...])
flatten(obj)	Flatten obj using this policy.
make_flatteners(myclass)	Return class-specific (flatten, unflatten) callables.
register(myclass)	Register myclass as a JAX pytree node and return the class.
unflatten(aux_data, children, myclass)	Reconstruct myclass from pytree data.

Attributes

ignore_defaults
ignore_keys
static_keys
static_types
validate_static

flatten(obj)

Flatten obj using this policy.

Return type:

tuple[tuple[Any, ...], tuple[tuple[str, ...], tuple[tuple[str, Any], ...]]]

make_flatteners(myclass)

Return class-specific (flatten, unflatten) callables.

Return type:

tuple[Callable[[Any], tuple[tuple[Any, ...], tuple[tuple[str, ...], tuple[tuple[str, Any], ...]]]], Callable[[tuple[tuple[str, ...], tuple[tuple[str, Any], ...]], tuple[Any, ...]], Any]]

register(myclass)

Register myclass as a JAX pytree node and return the class.

Returning the class lets this method be used either directly or as a decorator.

Return type:

type

unflatten(aux_data, children, myclass)

Reconstruct myclass from pytree data.

Return type:

Any