UniqueStateManager

class braintools.optim.UniqueStateManager(pytree=None)

Manage unique brainstate.State objects within a PyTree structure.

This class flattens a PyTree with State leaves to (path, leaf) pairs, removes duplicate State objects based on their id(), and supports recovering the unique states back to a PyTree structure.

Parameters:

pytree (State]) – If provided, make_unique() is called on it during construction.

Examples

>>> import jax.numpy as jnp
>>> from brainstate import ParamState
>>>
>>> # Create a pytree with some duplicate State objects
>>> state1 = ParamState(jnp.ones((2, 3)))
>>> state2 = ParamState(jnp.zeros((3, 4)))
>>> pytree = {
...     'layer1': {'weight': state1, 'bias': state2},
...     'layer2': {'weight': state1, 'bias': ParamState(jnp.ones((4,)))}  # duplicate
... }
>>>
>>> # Create manager and process the pytree
>>> manager = UniqueStateManager()
>>> unique_pytree = manager.make_unique(pytree)
>>> print(len(manager.flattened_states))  # 3, not 4 (duplicate removed)
3
>>>
>>> # Recover to pytree structure
>>> recovered = manager.to_pytree()

clear()

Clear all stored states and reset the manager.

get_flattened()

Get the flattened list of (path, state) pairs.

Returns:

List of (path, State) tuples for the unique states.

Return type:

List[Tuple[Any, State]]

get_state_by_path(target_path)

Retrieve a State object by its path.

Parameters:

target_path (Any) – The path to the desired State.

Returns:

The State object at the given path, or None if not found.

Return type:

State

make_unique(pytree)

Process a PyTree with State leaves and remove duplicates.

Parameters:

pytree (State]) – A PyTree whose leaves are brainstate.State objects.

Returns:

A PyTree containing only unique State objects (duplicates removed).

Return type:

State]

merge_with(other_pytree)

Merge another PyTree into the current unique states, maintaining uniqueness.

Parameters:

other_pytree (State]) – Another PyTree with State leaves to merge in.

Returns:

self, with the new unique states merged in.

Return type:

UniqueStateManager

property num_unique_states: int

Get the number of unique State objects.

to_dict()

Convert the stored unique states to a dict keyed by path strings.

Returns:

dict of {str – Dictionary whose keys are string representations of paths and whose values are the State objects.

Return type:

Dict[str, State]

to_dict_value()

Convert the stored unique states to a dict of path strings to State.value.

Returns:

dict of {str – Dictionary whose keys are string representations of paths and whose values are the corresponding State.value.

Return type:

Dict[str, Any]

to_pytree()

Convert the stored unique states back to a PyTree structure.

Returns:

PyTree with the unique State objects as leaves.

Return type:

State]

to_pytree_value()

Convert the stored unique states to a PyTree with State.value as leaves.

Returns:

PyTree where leaves are the values (State.value) of the State objects.

Return type:

PyTree

update_state(target_path, new_state)

Update a State object at a specific path.

Parameters:

• target_path (Any) – The path to the State to update.

• new_state (State) – The new State object.

Returns:

True if the update succeeded, False if the path was not found.

Return type:

bool