stringjax_tools.pytrees.flatten_func

Contents

stringjax_tools.pytrees.flatten_func#

flatten_func(obj, *, static_keys=(), ignore_keys=(), static_types=(<class 'str'>, <class 'bool'>), validate_static=True)#

Flatten obj for the JAX pytree protocol.

Parameters:

  • obj (Any) – Instance to flatten.

  • static_keys (Iterable[str]) – Attribute names stored as static auxiliary data.

  • ignore_keys (Iterable[str]) – Attribute names excluded from the pytree entirely. These attributes are absent after reconstruction unless the consuming package restores them in a custom unflatten wrapper. Do not put semantic configuration or user input here.

  • static_types (tuple[type, ...]) – Python value types that should always be static.

  • validate_static (bool) – If True, reject unhashable or array-like static values before they enter JAX auxiliary data.

Returns:

  • `` (children, aux_data)

  • ``jax.tree_util.register_pytree_node``. ``aux_data`` is structured as

  • `` (child_keys, static_items)

Return type:

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

Contents