brainstate 0.1.8__py2.py3-none-any.whl → 0.1.10__py2.py3-none-any.whl

brainstate/util/pretty_pytree.py

@@ -1,945 +1,945 @@
-# The file is adapted from the Flax library (https://github.com/google/flax).
-# The credit should go to the Flax authors.
-#
-# Copyright 2024 The Flax Authors.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from collections import abc
-from typing import TypeVar, Hashable, Union, Iterable, Any, Optional, Tuple, Dict
-
-import jax
-
-from brainstate.typing import Filter, PathParts
-from .pretty_repr import PrettyRepr, PrettyType, PrettyAttr, yield_unique_pretty_repr_items, pretty_repr_object
-from .struct import dataclass
-from .filter import to_predicate
-
-__all__ = [
-    'PrettyDict',
-    'NestedDict',
-    'FlattedDict',
-    'flat_mapping',
-    'nest_mapping',
-    'PrettyList',
-    'PrettyObject',
-]
-
-A = TypeVar('A')
-K = TypeVar('K', bound=Hashable)
-V = TypeVar('V')
-
-FlattedStateMapping = dict[PathParts, V]
-ExtractValueFn = abc.Callable[[Any], Any]
-SetValueFn = abc.Callable[[V, Any], V]
-
-
-class PrettyObject(PrettyRepr):
-    """
-    A class for generating a pretty representation of a tree-like structure.
-
-    This class extends the PrettyRepr class to provide a mechanism for
-    generating a human-readable, pretty representation of tree-like data
-    structures. It utilizes custom functions to represent the object and
-    its attributes in a structured and visually appealing format.
-
-    Methods
-    -------
-    __pretty_repr__: Generates a sequence of pretty representation items
-                     for the object.
-    __pretty_repr_item__: Returns a tuple of the key and value for pretty
-                          representation of an item in the data structure.
-    """
-
-    def __pretty_repr__(self):
-        """
-        Generates a pretty representation of the object.
-
-        This method yields a sequence of pretty representation items for the object,
-        using specified functions to represent the object and its attributes.
-
-        Yields:
-            Pretty representation items generated by `yield_unique_pretty_repr_items`.
-        """
-        yield from yield_unique_pretty_repr_items(
-            self,
-            repr_object=_repr_object_general,
-            repr_attr=_repr_attribute_general,
-        )
-
-    def __pretty_repr_item__(self, k, v):
-        """
-        Returns a tuple of the key and value for pretty representation.
-
-        This method is used to generate a pretty representation of an item
-        in a data structure, typically for debugging or logging purposes.
-
-        Args:
-            k: The key of the item.
-            v: The value of the item.
-
-        Returns:
-            A tuple containing the key and value.
-        """
-        return k, v
-
-
-PrettyReprTree = PrettyObject
-
-
-# the empty node is a struct.dataclass to be compatible with JAX.
-@dataclass
-class _EmptyNode:
-    pass
-
-
-_default_leaf = lambda *args: False
-empty_node = _EmptyNode()
-IsLeafCallable = abc.Callable[[Tuple[Any, ...], abc.Mapping[Any, Any]], bool]
-
-
-def flat_mapping(
-    xs: abc.Mapping[Any, Any],
-    /,
-    *,
-    keep_empty_nodes: bool = False,
-    is_leaf: Optional[IsLeafCallable] = _default_leaf,
-    sep: Optional[str] = None
-) -> 'FlattedDict':
-    """Flatten a nested mapping.
-
-    The nested keys are flattened to a tuple. See ``unflatten_mapping`` on how to
-    restore the nested mapping.
-
-    Example::
-
-      >>> xs = {'foo': 1, 'bar': {'a': 2, 'b': {}}}
-      >>> flat_xs = flat_mapping(xs)
-      >>> flat_xs
-      {('foo',): 1, ('bar', 'a'): 2}
-
-    Note that empty mappings are ignored and will not be restored by
-    ``unflatten_mapping``.
-
-    Args:
-      xs: A nested mapping.
-      keep_empty_nodes: replaces empty mappings with ``empty_node``.
-      is_leaf: An optional function that takes the next nested mapping and nested
-               keys and returns True if the nested mapping is a leaf (i.e., should not be
-               flattened further).
-      sep: If specified, then the keys of the returned mapping will be
-           ``sep``-joined strings (if ``None``, then keys will be tuples).
-
-    Returns:
-      The flattened mapping.
-    """
-    assert isinstance(xs, abc.Mapping), f'expected Mapping; got {type(xs).__qualname__}'
-
-    if sep is None:
-        def _key(path: Tuple[Any, ...]) -> Tuple[Any, ...] | str:
-            return path
-    else:
-
-        def _key(path: Tuple[Any, ...]) -> Tuple[Any, ...] | str:
-            return sep.join(path)
-
-    def _flatten(xs: Any, prefix: Tuple[Any, ...]) -> Dict[Any, Any]:
-        if not isinstance(xs, abc.Mapping) or is_leaf(prefix, xs):
-            return {_key(prefix): xs}
-
-        result = {}
-        is_empty = True
-        for key, value in xs.items():
-            is_empty = False
-            result.update(_flatten(value, prefix + (key,)))
-        if keep_empty_nodes and is_empty:
-            if prefix == ():  # when the whole input is empty
-                return {}
-            return {_key(prefix): empty_node}
-        return result
-
-    return FlattedDict(_flatten(xs, ()))
-
-
-def nest_mapping(
-    xs: Any,
-    /,
-    *,
-    sep: str | None = None
-) -> 'NestedDict':
-    """Unflatten a mapping.
-
-    See ``flatten_mapping``
-
-    Example::
-
-      >>> flat_xs = {
-      ...   ('foo',): 1,
-      ...   ('bar', 'a'): 2,
-      ... }
-      >>> xs = nest_mapping(flat_xs)
-      >>> xs
-      {'foo': 1, 'bar': {'a': 2}}
-
-    Args:
-      xs: a flattened mapping.
-      sep: separator (same as used with ``flatten_mapping()``).
-
-    Returns:
-      The nested mapping.
-    """
-    assert isinstance(xs, abc.Mapping), f'expected Mapping; got {type(xs).__qualname__}'
-    result: Dict[Any, Any] = {}
-    for path, value in xs.items():
-        if sep is not None:
-            path = path.split(sep)
-        if value is empty_node:
-            value = {}
-        cursor = result
-        for key in path[:-1]:
-            if key not in cursor:
-                cursor[key] = {}
-            cursor = cursor[key]
-        cursor[path[-1]] = value
-    return NestedDict(result)
-
-
-def _default_compare(x, values):
-    return id(x) in values
-
-
-def _default_process(x):
-    return id(x)
-
-
-class PrettyDict(dict, PrettyRepr):
-    __module__ = 'brainstate.util'
-
-    def __getattr__(self, key: K):  # type: ignore[misc]
-        return self[key]
-
-    def treefy_state(self):
-        """
-        Convert the :class:`State` objects to a reference tree of the state.
-        """
-        from brainstate._state import State
-        leaves, treedef = jax.tree.flatten(self)
-        leaves = jax.tree.map(lambda x: x.to_state_ref() if isinstance(x, State) else x, leaves)
-        return treedef.unflatten(leaves)
-
-    def to_dict(self) -> Dict[K, Dict[K, Any] | V]:
-        """
-        Convert the :class:`PrettyDict` to a dictionary.
-
-        Returns:
-          The dictionary.
-        """
-        return dict(self)  # type: ignore
-
-    def __repr__(self) -> str:
-        # repr the individual object with the pretty representation
-        return pretty_repr_object(self)
-
-    def __pretty_repr__(self):
-        yield from yield_unique_pretty_repr_items(self, _default_repr_object, _default_repr_attr)
-
-    def split(self, *filters) -> Union['PrettyDict[K, V]', Tuple['PrettyDict[K, V]', ...]]:
-        raise NotImplementedError
-
-    def filter(self, *filters) -> Union['PrettyDict[K, V]', Tuple['PrettyDict[K, V]', ...]]:
-        raise NotImplementedError
-
-    def merge(self, *states) -> 'PrettyDict[K, V]':
-        raise NotImplementedError
-
-    def subset(self, *filters) -> Union['PrettyDict[K, V]', Tuple['PrettyDict[K, V]', ...]]:
-        """
-        Subset a :class:`PrettyDict` into one or more :class:`PrettyDict`'s. The user must pass at least one
-        `:class:`Filter` (i.e. :class:`State`), and the filters must be exhaustive (i.e. they must cover all
-        :class:`State` types in the :class:`PrettyDict`).
-        """
-        return self.filter(*filters)
-
-
-class NestedStateRepr(PrettyRepr):
-    def __init__(self, state: PrettyDict):
-        self.state = state
-
-    def __pretty_repr__(self):
-        yield PrettyType('', value_sep=': ', start='{', end='}')
-
-        for r in self.state.__pretty_repr__():
-            if isinstance(r, PrettyType):
-                continue
-            yield r
-
-    def __treescope_repr__(self, path, subtree_renderer):
-        children = {}
-        for k, v in self.state.items():
-            if isinstance(v, PrettyDict):
-                v = NestedStateRepr(v)
-            children[k] = v
-        # Render as the dictionary itself at the same path.
-        return subtree_renderer(children, path=path)
-
-
-def _default_repr_object(node: PrettyDict):
-    yield PrettyType('', value_sep=': ', start='{', end='}')
-
-
-def _default_repr_attr(node):
-    for k, v in node.items():
-        if isinstance(v, list):
-            v = PrettyList(v)
-
-        if isinstance(v, dict):
-            v = PrettyDict(v)
-
-        if isinstance(v, PrettyDict):
-            v = NestedStateRepr(v)
-
-        yield PrettyAttr(repr(k), v)
-
-
-class NestedDict(PrettyDict):
-    """
-    A pytree-like structure that contains a :class:`Mapping` from strings or integers to leaves.
-
-    A valid leaf type is either :class:`State`, ``jax.Array``, ``numpy.ndarray`` or nested
-    :class:`NestedDict` and :class:`FlattedDict`.
-    """
-    __module__ = 'brainstate.util'
-
-    def __or__(self, other: 'NestedDict[K, V]') -> 'NestedDict[K, V]':
-        if not other:
-            return self
-        assert isinstance(other, NestedDict), f'expected NestedDict; got {type(other).__qualname__}'
-        return NestedDict.merge(self, other)
-
-    def __sub__(self, other: 'NestedDict[K, V]') -> 'NestedDict[K, V]':
-        if not other:
-            return self
-
-        assert isinstance(other, NestedDict), f'expected NestedDict; got {type(other).__qualname__}'
-        self_flat = self.to_flat()
-        other_flat = other.to_flat()
-        diff = {k: v for k, v in self_flat.items() if k not in other_flat}
-        return NestedDict.from_flat(diff)
-
-    def to_flat(self) -> 'FlattedDict':
-        """
-        Flatten the nested mapping into a flat mapping.
-
-        Returns:
-            The flattened mapping.
-        """
-        return flat_mapping(self)
-
-    @classmethod
-    def from_flat(cls, flat_dict: abc.Mapping[PathParts, V] | Iterable[tuple[PathParts, V]]) -> 'NestedDict':
-        """
-        Create a :class:`NestedDict` from a flat mapping.
-
-        Args:
-          flat_dict: The flat mapping.
-
-        Returns:
-          The :class:`NestedDict`.
-        """
-        nested_state = nest_mapping(dict(flat_dict))
-        return cls(nested_state)
-
-    def split(  # type: ignore[misc]
-        self,
-        first: Filter,
-        /,
-        *filters: Filter
-    ) -> Union['NestedDict[K, V]', Tuple['NestedDict[K, V]', ...]]:
-        """
-        Split a :class:`NestedDict` into one or more :class:`NestedDict`'s. The
-        user must pass at least one `:class:`Filter` (i.e. :class:`State`),
-        and the filters must be exhaustive (i.e. they must cover all
-        :class:`State` types in the :class:`NestedDict`).
-
-        Example usage::
-
-          >>> import brainstate as brainstate
-
-          >>> class Model(brainstate.nn.Module):
-          ...   def __init__(self):
-          ...     super().__init__()
-          ...     self.batchnorm = brainstate.nn.BatchNorm1d([10, 3])
-          ...     self.linear = brainstate.nn.Linear(2, 3)
-          ...   def __call__(self, x):
-          ...     return self.linear(self.batchnorm(x))
-
-          >>> model = Model()
-          >>> state_map = brainstate.graph.treefy_states(model)
-          >>> param, others = state_map.treefy_split(brainstate.ParamState, ...)
-
-        Arguments:
-          first: The first filter
-          *filters: The optional, additional filters to group the state into mutually exclusive substates.
-
-        Returns:
-          One or more ``States`` equal to the number of filters passed.
-        """
-        filters = (first, *filters)
-        *states_, rest = _split_nested_mapping(self, *filters)
-        if rest:
-            raise ValueError(f'Non-exhaustive filters, got a non-empty remainder: {rest}.\n'
-                             f'Use `...` to match all remaining elements.')
-
-        states: NestedDict | Tuple[NestedDict, ...]
-        if len(states_) == 1:
-            states = states_[0]
-        else:
-            states = tuple(states_)
-        return states  # type: ignore[bad-return-type]
-
-    def filter(
-        self,
-        first: Filter,
-        /,
-        *filters: Filter,
-    ) -> Union['NestedDict[K, V]', Tuple['NestedDict[K, V]', ...]]:
-        """
-        Filter a :class:`NestedDict` into one or more :class:`NestedDict`'s. The
-        user must pass at least one `:class:`Filter` (i.e. :class:`State`).
-        This method is similar to :meth:`split() <flax.nnx.NestedDict.state.split>`,
-        except the filters can be non-exhaustive.
-
-        Arguments:
-          first: The first filter
-          *filters: The optional, additional filters to group the state into mutually exclusive substates.
-
-        Returns:
-          One or more ``States`` equal to the number of filters passed.
-        """
-        *states_, _rest = _split_nested_mapping(self, first, *filters)
-        assert len(states_) == len(filters) + 1, f'Expected {len(filters) + 1} states, got {len(states_)}'
-        if len(states_) == 1:
-            states = states_[0]
-        else:
-            states = tuple(states_)
-        return states  # type: ignore[bad-return-type]
-
-    @staticmethod
-    def merge(
-        state: Union['NestedDict[K, V]', 'FlattedDict[K, V]'],
-        /,
-        *states: Union['NestedDict[K, V]', 'FlattedDict[K, V]']
-    ) -> 'NestedDict[K, V]':
-        """
-        The inverse of :meth:`split()`.
-
-        ``merge`` takes one or more :class:`PrettyDict`'s and creates a new :class:`PrettyDict`.
-
-        Args:
-          state: A :class:`PrettyDict` object.
-          *states: Additional :class:`PrettyDict` objects.
-
-        Returns:
-          The merged :class:`PrettyDict`.
-        """
-        if not states:
-            return state
-        states = (state, *states)
-        new_state: FlattedDict = FlattedDict()
-        for state in states:
-            if isinstance(state, NestedDict):
-                new_state.update(state.to_flat())  # type: ignore[attribute-error] # pytype is wrong here
-            elif isinstance(state, FlattedDict):
-                new_state.update(state)
-            else:
-                raise TypeError(f'Expected Nested or Flatted Mapping, got {type(state)} instead.')
-        return NestedDict.from_flat(new_state)
-
-    def to_pure_dict(self) -> Dict[str, Any]:
-        flat_values = {k: x for k, x in self.to_flat().items()}
-        return nest_mapping(flat_values).to_dict()
-
-    def replace_by_pure_dict(
-        self,
-        pure_dict: Dict[str, Any],
-        replace_fn: Optional[SetValueFn] = None
-    ):
-        if replace_fn is None:
-            replace_fn = lambda x, v: x.replace(v) if hasattr(x, 'replace') else v
-        current_flat = self.to_flat()
-        for kp, v in flat_mapping(pure_dict).items():
-            if kp not in current_flat:
-                raise ValueError(f'key in pure_dict not available in state: {kp}')
-            current_flat[kp] = replace_fn(current_flat[kp], v)
-        self.update(nest_mapping(current_flat))
-
-
-class FlattedDict(PrettyDict):
-    """
-    A pytree-like structure that contains a :class:`Mapping` from strings or integers to leaves.
-
-    A valid leaf type is either :class:`State`, ``jax.Array``, ``numpy.ndarray`` or Python variables.
-
-    A :class:`NestedDict` can be generated by either calling :func:`states()` or
-    :func:`nodes()` on the :class:`Module`.
-
-    Example usage::
-
-      >>> import brainstate as brainstate
-      >>> import jax.numpy as jnp
-      >>>
-      >>> class Model(brainstate.nn.Module):
-      ...   def __init__(self):
-      ...     super().__init__()
-      ...     self.batchnorm = brainstate.nn.BatchNorm1d([10, 3])
-      ...     self.linear = brainstate.nn.Linear(2, 3)
-      ...   def __call__(self, x):
-      ...     return self.linear(self.batchnorm(x))
-      >>>
-      >>> model = Model()
-
-      >>> # retrieve the states of the model
-      >>> model.states()  # with the same to the function of ``brainstate.graph.states()``
-      FlattedDict({
-        ('batchnorm', 'running_mean'): LongTermState(
-          value=Array([[0., 0., 0.]], dtype=float32)
-        ),
-        ('batchnorm', 'running_var'): LongTermState(
-          value=Array([[1., 1., 1.]], dtype=float32)
-        ),
-        ('batchnorm', 'weight'): ParamState(
-          value={'bias': Array([[0., 0., 0.]], dtype=float32), 'scale': Array([[1., 1., 1.]], dtype=float32)}
-        ),
-        ('linear', 'weight'): ParamState(
-          value={'weight': Array([[-0.21467684,  0.7621282 , -0.50756454, -0.49047297],
-                 [-0.90413696,  0.6711    , -0.1254792 ,  0.50412565],
-                 [ 0.23975602,  0.47905368,  1.4851435 ,  0.16745673]],      dtype=float32), 'bias': Array([0., 0., 0., 0.], dtype=float32)}
-        )
-      })
-
-      >>> # retrieve the nodes of the model
-      >>> model.nodes()  # with the same to the function of ``brainstate.graph.nodes()``
-      FlattedDict({
-        ('batchnorm',): BatchNorm1d(
-          in_size=(10, 3),
-          out_size=(10, 3),
-          affine=True,
-          bias_initializer=Constant(value=0.0, dtype=<class 'numpy.float32'>),
-          scale_initializer=Constant(value=1.0, dtype=<class 'numpy.float32'>),
-          dtype=<class 'numpy.float32'>,
-          track_running_stats=True,
-          momentum=Array(shape=(), dtype=float32),
-          epsilon=Array(shape=(), dtype=float32),
-          feature_axis=(1,),
-          axis_name=None,
-          axis_index_groups=None,
-          running_mean=LongTermState(
-            value=Array(shape=(1, 3), dtype=float32)
-          ),
-          running_var=LongTermState(
-            value=Array(shape=(1, 3), dtype=float32)
-          ),
-          weight=ParamState(
-            value={'bias': Array(shape=(1, 3), dtype=float32), 'scale': Array(shape=(1, 3), dtype=float32)}
-          )
-        ),
-        ('linear',): Linear(
-          in_size=(10, 3),
-          out_size=(10, 4),
-          w_mask=None,
-          weight=ParamState(
-            value={'bias': Array(shape=(4,), dtype=float32), 'weight': Array(shape=(3, 4), dtype=float32)}
-          )
-        ),
-        (): Model(
-          batchnorm=BatchNorm1d(...),
-          linear=Linear(...)
-        )
-      })
-    """
-    __module__ = 'brainstate.util'
-
-    def __or__(self, other: 'FlattedDict[K, V]') -> 'FlattedDict[K, V]':
-        if not other:
-            return self
-        assert isinstance(other, FlattedDict), f'expected NestedDict; got {type(other).__qualname__}'
-        return FlattedDict.merge(self, other)
-
-    def __sub__(self, other: 'FlattedDict[K, V]') -> 'FlattedDict[K, V]':
-        if not other:
-            return self
-        assert isinstance(other, FlattedDict), f'expected NestedDict; got {type(other).__qualname__}'
-        diff = {k: v for k, v in self.items() if k not in other}
-        return FlattedDict(diff)
-
-    def to_nest(self) -> NestedDict:
-        """
-        Unflatten the flat mapping into a nested mapping.
-
-        Returns:
-            The nested mapping.
-        """
-        return nest_mapping(self)
-
-    @classmethod
-    def from_nest(
-        cls, nested_dict: abc.Mapping[PathParts, V] | Iterable[tuple[PathParts, V]],
-    ) -> 'FlattedDict':
-        """
-        Create a :class:`NestedDict` from a flat mapping.
-
-        Args:
-          nested_dict: The flat mapping.
-
-        Returns:
-          The :class:`NestedDict`.
-        """
-        return flat_mapping(nested_dict)
-
-    def split(  # type: ignore[misc]
-        self,
-        first: Filter,
-        /,
-        *filters: Filter
-    ) -> Union['FlattedDict[K, V]', tuple['FlattedDict[K, V]', ...]]:
-        """
-        Split a :class:`FlattedDict` into one or more :class:`FlattedDict`'s. The
-        user must pass at least one `:class:`Filter` (i.e. :class:`State`),
-        and the filters must be exhaustive (i.e. they must cover all
-        :class:`State` types in the :class:`NestedDict`).
-
-        Arguments:
-            first: The first filter
-            *filters: The optional, additional filters to group the state into mutually exclusive substates.
-
-        Returns:
-            One or more ``States`` equal to the number of filters passed.
-        """
-        filters = (first, *filters)
-        *states_, rest = _split_flatted_mapping(self, *filters)
-        if rest:
-            raise ValueError(f'Non-exhaustive filters, got a non-empty remainder: {rest}.\n'
-                             f'Use `...` to match all remaining elements.')
-
-        states: FlattedDict | Tuple[FlattedDict, ...]
-        if len(states_) == 1:
-            states = states_[0]
-        else:
-            states = tuple(states_)
-        return states  # type: ignore[bad-return-type]
-
-    def filter(
-        self,
-        first: Filter,
-        /,
-        *filters: Filter,
-    ) -> Union['FlattedDict[K, V]', Tuple['FlattedDict[K, V]', ...]]:
-        """
-        Filter a :class:`FlattedDict` into one or more :class:`FlattedDict`'s. The
-        user must pass at least one `:class:`Filter` (i.e. :class:`State`).
-        This method is similar to :meth:`split() <flax.nnx.NestedDict.state.split>`,
-        except the filters can be non-exhaustive.
-
-        Arguments:
-            first: The first filter
-            *filters: The optional, additional filters to group the state into mutually exclusive substates.
-
-        Returns:
-            One or more ``States`` equal to the number of filters passed.
-        """
-        *states_, _rest = _split_flatted_mapping(self, first, *filters)
-        assert len(states_) == len(filters) + 1, f'Expected {len(filters) + 1} states, got {len(states_)}'
-        if len(states_) == 1:
-            states = states_[0]
-        else:
-            states = tuple(states_)
-        return states  # type: ignore[bad-return-type]
-
-    @staticmethod
-    def merge(
-        state: Union['FlattedDict[K, V]', 'NestedDict[K, V]'],
-        /,
-        *states: Union['FlattedDict[K, V]', 'NestedDict[K, V]']
-    ) -> 'FlattedDict[K, V]':
-        """
-        The inverse of :meth:`split()`.
-
-        ``merge`` takes one or more :class:`FlattedDict`'s and creates a new :class:`FlattedDict`.
-
-        Args:
-          state: A :class:`PrettyDict` object.
-          *states: Additional :class:`PrettyDict` objects.
-
-        Returns:
-          The merged :class:`PrettyDict`.
-        """
-        if not states:
-            return state
-        states = (state, *states)
-        new_state: FlattedStateMapping[V] = {}
-        for state in states:
-            if isinstance(state, NestedDict):
-                new_state.update(state.to_flat())  # type: ignore[attribute-error] # pytype is wrong here
-            elif isinstance(state, FlattedDict):
-                new_state.update(state)
-            else:
-                raise TypeError(f'Expected Nested or Flatted Mapping, got {type(state)} instead.')
-        return FlattedDict(new_state)
-
-    def to_dict_values(self):
-        """
-        Convert a FlattedDict containing State objects to a plain dictionary of values.
-        
-        This method extracts the underlying values from any State objects in the FlattedDict,
-        creating a new dictionary with the same keys but where each State object is replaced
-        by its value attribute. Non-State objects are kept as is.
-        
-        Returns:
-            dict: A dictionary with the same keys as the FlattedDict, but where each State
-                  object is replaced by its value attribute. Non-State objects remain unchanged.
-        
-        Example:
-            >>> flat_dict = FlattedDict({('model', 'layer1', 'weight'): ParamState(value=jnp.ones((10, 5)))})
-            >>> flat_dict.to_dict_values()
-            {('model', 'layer1', 'weight'): Array([[1., 1., ...]], dtype=float32)}
-        """
-        from brainstate._state import State
-        return {
-            k: v.value if isinstance(v, State) else v
-            for k, v in self.items()
-        }
-
-    def assign_dict_values(self, data: dict):
-        """
-        Assign values from a dictionary to this FlattedDict.
-        
-        This method updates the values in the FlattedDict with values from the provided
-        dictionary. For keys that correspond to State objects, the value attribute of
-        the State is updated. For other keys, the value in the FlattedDict is directly
-        replaced with the new value.
-        
-        The method requires that all keys in the FlattedDict exist in the provided
-        dictionary, otherwise a KeyError is raised.
-        
-        Args:
-            data (dict): A dictionary containing the values to assign, where keys 
-                         must match those in the FlattedDict.
-        
-        Raises:
-            KeyError: If a key in the FlattedDict is not present in the provided dictionary.
-        
-        Example:
-            >>> flat_dict = FlattedDict({('model', 'weight'): ParamState(value=jnp.zeros((5, 5)))})
-            >>> flat_dict.assign_dict_values({('model', 'weight'): jnp.ones((5, 5))})
-            # The ParamState's value is now an array of ones
-        """
-        from brainstate._state import State
-        for k in self.keys():
-            if k not in data:
-                raise KeyError(f'Invalid key: {k!r}')
-            val = self[k]
-            if isinstance(val, State):
-                val.value = data[k]
-            else:
-                self[k] = data[k]
-
-
-def _split_nested_mapping(
-    mapping: 'NestedDict[K, V]',
-    *filters: Filter,
-) -> Tuple['NestedDict[K, V]', ...]:
-    # check if the filters are exhaustive
-    for i, filter_ in enumerate(filters):
-        if filter_ in (..., True) and i != len(filters) - 1:
-            remaining_filters = filters[i + 1:]
-            if not all(f in (..., True) for f in remaining_filters):
-                raise ValueError('`...` or `True` can only be used as the last filters, '
-                                 f'got {filter_} it at index {i}.')
-
-    # change the filters to predicates
-    predicates = tuple(map(to_predicate, filters))
-
-    # we have n + 1 state mappings, where n is the number of predicates
-    # the last state mapping is for values that don't match any predicate
-    flat_states: tuple[FlattedStateMapping[V], ...] = tuple({} for _ in range(len(predicates) + 1))
-
-    assert isinstance(mapping, NestedDict), f'expected NestedDict; got {type(mapping).__qualname__}'
-    flat_state = mapping.to_flat()
-    for path, value in flat_state.items():
-        for i, predicate in enumerate(predicates):
-            if predicate(path, value):
-                flat_states[i][path] = value  # type: ignore[index] # mypy is wrong here?
-                break
-        else:
-            # if we didn't break, set leaf to last state
-            flat_states[-1][path] = value  # type: ignore[index] # mypy is wrong here?
-
-    return tuple(NestedDict.from_flat(flat_state) for flat_state in flat_states)
-
-
-def _split_flatted_mapping(
-    mapping: FlattedDict[K, V],
-    *filters: Filter,
-) -> Tuple[FlattedDict[K, V], ...]:
-    # check if the filters are exhaustive
-    for i, filter_ in enumerate(filters):
-        if filter_ in (..., True) and i != len(filters) - 1:
-            remaining_filters = filters[i + 1:]
-            if not all(f in (..., True) for f in remaining_filters):
-                raise ValueError('`...` or `True` can only be used as the last filters, '
-                                 f'got {filter_} it at index {i}.')
-
-    # change the filters to predicates
-    predicates = tuple(map(to_predicate, filters))
-
-    # we have n + 1 state mappings, where n is the number of predicates
-    # the last state mapping is for values that don't match any predicate
-    flat_states: tuple[FlattedStateMapping[V], ...] = tuple({} for _ in range(len(predicates) + 1))
-
-    assert isinstance(mapping, FlattedDict), f'expected FlattedDict; got {type(mapping).__qualname__}'
-    for path, value in mapping.items():
-        for i, predicate in enumerate(predicates):
-            if predicate(path, value):
-                flat_states[i][path] = value  # type: ignore[index] # mypy is wrong here?
-                break
-        else:
-            # if we didn't break, set leaf to last state
-            flat_states[-1][path] = value  # type: ignore[index] # mypy is wrong here?
-
-    return tuple(FlattedDict(flat_state) for flat_state in flat_states)
-
-
-# register :class:`NestedDict` as a pytree
-def _nest_flatten_with_keys(x: NestedDict):
-    items = sorted(x.items())
-    children = tuple((jax.tree_util.DictKey(key), value) for key, value in items)
-    return children, tuple(key for key, _ in items)
-
-
-def _nest_unflatten(
-    static: Tuple[K, ...],
-    leaves: Tuple[V, ...] | Tuple[Dict[K, V]],
-):
-    return NestedDict(zip(static, leaves))
-
-
-jax.tree_util.register_pytree_with_keys(NestedDict,
-                                        _nest_flatten_with_keys,
-                                        _nest_unflatten)  # type: ignore[arg-type]
-
-
-# register :class:`FlattedDict` as a pytree
-
-def _flat_unflatten(
-    static: Tuple[K, ...],
-    leaves: Tuple[V, ...] | Tuple[Dict[K, V]],
-):
-    return FlattedDict(zip(static, leaves))
-
-
-jax.tree_util.register_pytree_with_keys(FlattedDict,
-                                        _nest_flatten_with_keys,
-                                        _flat_unflatten)  # type: ignore[arg-type]
-
-
-@jax.tree_util.register_pytree_node_class
-class PrettyList(list, PrettyRepr):
-    __module__ = 'brainstate.util'
-
-    def __pretty_repr__(self):
-        yield from yield_unique_pretty_repr_items(self, _list_repr_object, _list_repr_attr)
-
-    def __repr__(self):
-        return pretty_repr_object(self)
-
-    def tree_flatten(self):
-        return list(self), ()
-
-    @classmethod
-    def tree_unflatten(cls, aux_data, children):
-        return cls(children)
-
-
-def _list_repr_attr(node: PrettyList):
-    for v in node:
-        if isinstance(v, list):
-            v = PrettyList(v)
-        if isinstance(v, dict):
-            v = PrettyDict(v)
-        if isinstance(v, PrettyDict):
-            v = NestedStateRepr(v)
-        yield PrettyAttr('', v)
-
-
-def _list_repr_object(node: PrettyDict):
-    yield PrettyType('', value_sep='', start='[', end=']')
-
-
-def _repr_object_general(node: PrettyDict):
-    """
-    Generate a general representation of a PrettyDict object.
-
-    This function is used to create a pretty representation of a PrettyDict
-    object, which includes the type of the object and its value separator.
-
-    Args:
-        node (PrettyDict): The PrettyDict object to be represented.
-
-    Yields:
-        PrettyType: A PrettyType object representing the type of the node,
-        with specified value separator, start, and end characters.
-    """
-    yield PrettyType(type(node), value_sep='=', start='(', end=')')
-
-
-def _repr_attribute_general(node):
-    """
-    Generate a pretty representation of the attributes of a node.
-
-    This function iterates over the attributes of a given node and attempts
-    to generate a pretty representation for each attribute. It handles
-    conversion of lists and dictionaries to their pretty representation
-    counterparts and yields a PrettyAttr object for each attribute.
-
-    Args:
-        node: The object whose attributes are to be represented.
-
-    Yields:
-        PrettyAttr: A PrettyAttr object representing the key and value of
-        each attribute in a pretty format.
-    """
-    for k, v in vars(node).items():
-        try:
-            res = node.__pretty_repr_item__(k, v)
-            if res is None:
-                continue
-            k, v = res
-        except AttributeError:
-            pass
-
-        if k is None:
-            continue
-
-        # convert list to PrettyList
-        if isinstance(v, list):
-            v = PrettyList(v)
-
-        # convert dict to PrettyDict
-        if isinstance(v, dict):
-            v = PrettyDict(v)
-
-        # convert PrettyDict to NestedStateRepr
-        if isinstance(v, PrettyDict):
-            v = NestedStateRepr(v)
-
-        yield PrettyAttr(k, v)
+# The file is adapted from the Flax library (https://github.com/google/flax).
+# The credit should go to the Flax authors.
+#
+# Copyright 2024 The Flax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from collections import abc
+from typing import TypeVar, Hashable, Union, Iterable, Any, Optional, Tuple, Dict
+
+import jax
+
+from brainstate.typing import Filter, PathParts
+from .pretty_repr import PrettyRepr, PrettyType, PrettyAttr, yield_unique_pretty_repr_items, pretty_repr_object
+from .struct import dataclass
+from .filter import to_predicate
+
+__all__ = [
+    'PrettyDict',
+    'NestedDict',
+    'FlattedDict',
+    'flat_mapping',
+    'nest_mapping',
+    'PrettyList',
+    'PrettyObject',
+]
+
+A = TypeVar('A')
+K = TypeVar('K', bound=Hashable)
+V = TypeVar('V')
+
+FlattedStateMapping = dict[PathParts, V]
+ExtractValueFn = abc.Callable[[Any], Any]
+SetValueFn = abc.Callable[[V, Any], V]
+
+
+class PrettyObject(PrettyRepr):
+    """
+    A class for generating a pretty representation of a tree-like structure.
+
+    This class extends the PrettyRepr class to provide a mechanism for
+    generating a human-readable, pretty representation of tree-like data
+    structures. It utilizes custom functions to represent the object and
+    its attributes in a structured and visually appealing format.
+
+    Methods
+    -------
+    __pretty_repr__: Generates a sequence of pretty representation items
+                     for the object.
+    __pretty_repr_item__: Returns a tuple of the key and value for pretty
+                          representation of an item in the data structure.
+    """
+
+    def __pretty_repr__(self):
+        """
+        Generates a pretty representation of the object.
+
+        This method yields a sequence of pretty representation items for the object,
+        using specified functions to represent the object and its attributes.
+
+        Yields:
+            Pretty representation items generated by `yield_unique_pretty_repr_items`.
+        """
+        yield from yield_unique_pretty_repr_items(
+            self,
+            repr_object=_repr_object_general,
+            repr_attr=_repr_attribute_general,
+        )
+
+    def __pretty_repr_item__(self, k, v):
+        """
+        Returns a tuple of the key and value for pretty representation.
+
+        This method is used to generate a pretty representation of an item
+        in a data structure, typically for debugging or logging purposes.
+
+        Args:
+            k: The key of the item.
+            v: The value of the item.
+
+        Returns:
+            A tuple containing the key and value.
+        """
+        return k, v
+
+
+PrettyReprTree = PrettyObject
+
+
+# the empty node is a struct.dataclass to be compatible with JAX.
+@dataclass
+class _EmptyNode:
+    pass
+
+
+_default_leaf = lambda *args: False
+empty_node = _EmptyNode()
+IsLeafCallable = abc.Callable[[Tuple[Any, ...], abc.Mapping[Any, Any]], bool]
+
+
+def flat_mapping(
+    xs: abc.Mapping[Any, Any],
+    /,
+    *,
+    keep_empty_nodes: bool = False,
+    is_leaf: Optional[IsLeafCallable] = _default_leaf,
+    sep: Optional[str] = None
+) -> 'FlattedDict':
+    """Flatten a nested mapping.
+
+    The nested keys are flattened to a tuple. See ``unflatten_mapping`` on how to
+    restore the nested mapping.
+
+    Example::
+
+      >>> xs = {'foo': 1, 'bar': {'a': 2, 'b': {}}}
+      >>> flat_xs = flat_mapping(xs)
+      >>> flat_xs
+      {('foo',): 1, ('bar', 'a'): 2}
+
+    Note that empty mappings are ignored and will not be restored by
+    ``unflatten_mapping``.
+
+    Args:
+      xs: A nested mapping.
+      keep_empty_nodes: replaces empty mappings with ``empty_node``.
+      is_leaf: An optional function that takes the next nested mapping and nested
+               keys and returns True if the nested mapping is a leaf (i.e., should not be
+               flattened further).
+      sep: If specified, then the keys of the returned mapping will be
+           ``sep``-joined strings (if ``None``, then keys will be tuples).
+
+    Returns:
+      The flattened mapping.
+    """
+    assert isinstance(xs, abc.Mapping), f'expected Mapping; got {type(xs).__qualname__}'
+
+    if sep is None:
+        def _key(path: Tuple[Any, ...]) -> Tuple[Any, ...] | str:
+            return path
+    else:
+
+        def _key(path: Tuple[Any, ...]) -> Tuple[Any, ...] | str:
+            return sep.join(path)
+
+    def _flatten(xs: Any, prefix: Tuple[Any, ...]) -> Dict[Any, Any]:
+        if not isinstance(xs, abc.Mapping) or is_leaf(prefix, xs):
+            return {_key(prefix): xs}
+
+        result = {}
+        is_empty = True
+        for key, value in xs.items():
+            is_empty = False
+            result.update(_flatten(value, prefix + (key,)))
+        if keep_empty_nodes and is_empty:
+            if prefix == ():  # when the whole input is empty
+                return {}
+            return {_key(prefix): empty_node}
+        return result
+
+    return FlattedDict(_flatten(xs, ()))
+
+
+def nest_mapping(
+    xs: Any,
+    /,
+    *,
+    sep: str | None = None
+) -> 'NestedDict':
+    """Unflatten a mapping.
+
+    See ``flatten_mapping``
+
+    Example::
+
+      >>> flat_xs = {
+      ...   ('foo',): 1,
+      ...   ('bar', 'a'): 2,
+      ... }
+      >>> xs = nest_mapping(flat_xs)
+      >>> xs
+      {'foo': 1, 'bar': {'a': 2}}
+
+    Args:
+      xs: a flattened mapping.
+      sep: separator (same as used with ``flatten_mapping()``).
+
+    Returns:
+      The nested mapping.
+    """
+    assert isinstance(xs, abc.Mapping), f'expected Mapping; got {type(xs).__qualname__}'
+    result: Dict[Any, Any] = {}
+    for path, value in xs.items():
+        if sep is not None:
+            path = path.split(sep)
+        if value is empty_node:
+            value = {}
+        cursor = result
+        for key in path[:-1]:
+            if key not in cursor:
+                cursor[key] = {}
+            cursor = cursor[key]
+        cursor[path[-1]] = value
+    return NestedDict(result)
+
+
+def _default_compare(x, values):
+    return id(x) in values
+
+
+def _default_process(x):
+    return id(x)
+
+
+class PrettyDict(dict, PrettyRepr):
+    __module__ = 'brainstate.util'
+
+    def __getattr__(self, key: K):  # type: ignore[misc]
+        return self[key]
+
+    def treefy_state(self):
+        """
+        Convert the :class:`State` objects to a reference tree of the state.
+        """
+        from brainstate._state import State
+        leaves, treedef = jax.tree.flatten(self)
+        leaves = jax.tree.map(lambda x: x.to_state_ref() if isinstance(x, State) else x, leaves)
+        return treedef.unflatten(leaves)
+
+    def to_dict(self) -> Dict[K, Dict[K, Any] | V]:
+        """
+        Convert the :class:`PrettyDict` to a dictionary.
+
+        Returns:
+          The dictionary.
+        """
+        return dict(self)  # type: ignore
+
+    def __repr__(self) -> str:
+        # repr the individual object with the pretty representation
+        return pretty_repr_object(self)
+
+    def __pretty_repr__(self):
+        yield from yield_unique_pretty_repr_items(self, _default_repr_object, _default_repr_attr)
+
+    def split(self, *filters) -> Union['PrettyDict[K, V]', Tuple['PrettyDict[K, V]', ...]]:
+        raise NotImplementedError
+
+    def filter(self, *filters) -> Union['PrettyDict[K, V]', Tuple['PrettyDict[K, V]', ...]]:
+        raise NotImplementedError
+
+    def merge(self, *states) -> 'PrettyDict[K, V]':
+        raise NotImplementedError
+
+    def subset(self, *filters) -> Union['PrettyDict[K, V]', Tuple['PrettyDict[K, V]', ...]]:
+        """
+        Subset a :class:`PrettyDict` into one or more :class:`PrettyDict`'s. The user must pass at least one
+        `:class:`Filter` (i.e. :class:`State`), and the filters must be exhaustive (i.e. they must cover all
+        :class:`State` types in the :class:`PrettyDict`).
+        """
+        return self.filter(*filters)
+
+
+class NestedStateRepr(PrettyRepr):
+    def __init__(self, state: PrettyDict):
+        self.state = state
+
+    def __pretty_repr__(self):
+        yield PrettyType('', value_sep=': ', start='{', end='}')
+
+        for r in self.state.__pretty_repr__():
+            if isinstance(r, PrettyType):
+                continue
+            yield r
+
+    def __treescope_repr__(self, path, subtree_renderer):
+        children = {}
+        for k, v in self.state.items():
+            if isinstance(v, PrettyDict):
+                v = NestedStateRepr(v)
+            children[k] = v
+        # Render as the dictionary itself at the same path.
+        return subtree_renderer(children, path=path)
+
+
+def _default_repr_object(node: PrettyDict):
+    yield PrettyType('', value_sep=': ', start='{', end='}')
+
+
+def _default_repr_attr(node):
+    for k, v in node.items():
+        if isinstance(v, list):
+            v = PrettyList(v)
+
+        if isinstance(v, dict):
+            v = PrettyDict(v)
+
+        if isinstance(v, PrettyDict):
+            v = NestedStateRepr(v)
+
+        yield PrettyAttr(repr(k), v)
+
+
+class NestedDict(PrettyDict):
+    """
+    A pytree-like structure that contains a :class:`Mapping` from strings or integers to leaves.
+
+    A valid leaf type is either :class:`State`, ``jax.Array``, ``numpy.ndarray`` or nested
+    :class:`NestedDict` and :class:`FlattedDict`.
+    """
+    __module__ = 'brainstate.util'
+
+    def __or__(self, other: 'NestedDict[K, V]') -> 'NestedDict[K, V]':
+        if not other:
+            return self
+        assert isinstance(other, NestedDict), f'expected NestedDict; got {type(other).__qualname__}'
+        return NestedDict.merge(self, other)
+
+    def __sub__(self, other: 'NestedDict[K, V]') -> 'NestedDict[K, V]':
+        if not other:
+            return self
+
+        assert isinstance(other, NestedDict), f'expected NestedDict; got {type(other).__qualname__}'
+        self_flat = self.to_flat()
+        other_flat = other.to_flat()
+        diff = {k: v for k, v in self_flat.items() if k not in other_flat}
+        return NestedDict.from_flat(diff)
+
+    def to_flat(self) -> 'FlattedDict':
+        """
+        Flatten the nested mapping into a flat mapping.
+
+        Returns:
+            The flattened mapping.
+        """
+        return flat_mapping(self)
+
+    @classmethod
+    def from_flat(cls, flat_dict: abc.Mapping[PathParts, V] | Iterable[tuple[PathParts, V]]) -> 'NestedDict':
+        """
+        Create a :class:`NestedDict` from a flat mapping.
+
+        Args:
+          flat_dict: The flat mapping.
+
+        Returns:
+          The :class:`NestedDict`.
+        """
+        nested_state = nest_mapping(dict(flat_dict))
+        return cls(nested_state)
+
+    def split(  # type: ignore[misc]
+        self,
+        first: Filter,
+        /,
+        *filters: Filter
+    ) -> Union['NestedDict[K, V]', Tuple['NestedDict[K, V]', ...]]:
+        """
+        Split a :class:`NestedDict` into one or more :class:`NestedDict`'s. The
+        user must pass at least one `:class:`Filter` (i.e. :class:`State`),
+        and the filters must be exhaustive (i.e. they must cover all
+        :class:`State` types in the :class:`NestedDict`).
+
+        Example usage::
+
+          >>> import brainstate as brainstate
+
+          >>> class Model(brainstate.nn.Module):
+          ...   def __init__(self):
+          ...     super().__init__()
+          ...     self.batchnorm = brainstate.nn.BatchNorm1d([10, 3])
+          ...     self.linear = brainstate.nn.Linear(2, 3)
+          ...   def __call__(self, x):
+          ...     return self.linear(self.batchnorm(x))
+
+          >>> model = Model()
+          >>> state_map = brainstate.graph.treefy_states(model)
+          >>> param, others = state_map.treefy_split(brainstate.ParamState, ...)
+
+        Arguments:
+          first: The first filter
+          *filters: The optional, additional filters to group the state into mutually exclusive substates.
+
+        Returns:
+          One or more ``States`` equal to the number of filters passed.
+        """
+        filters = (first, *filters)
+        *states_, rest = _split_nested_mapping(self, *filters)
+        if rest:
+            raise ValueError(f'Non-exhaustive filters, got a non-empty remainder: {rest}.\n'
+                             f'Use `...` to match all remaining elements.')
+
+        states: NestedDict | Tuple[NestedDict, ...]
+        if len(states_) == 1:
+            states = states_[0]
+        else:
+            states = tuple(states_)
+        return states  # type: ignore[bad-return-type]
+
+    def filter(
+        self,
+        first: Filter,
+        /,
+        *filters: Filter,
+    ) -> Union['NestedDict[K, V]', Tuple['NestedDict[K, V]', ...]]:
+        """
+        Filter a :class:`NestedDict` into one or more :class:`NestedDict`'s. The
+        user must pass at least one `:class:`Filter` (i.e. :class:`State`).
+        This method is similar to :meth:`split() <flax.nnx.NestedDict.state.split>`,
+        except the filters can be non-exhaustive.
+
+        Arguments:
+          first: The first filter
+          *filters: The optional, additional filters to group the state into mutually exclusive substates.
+
+        Returns:
+          One or more ``States`` equal to the number of filters passed.
+        """
+        *states_, _rest = _split_nested_mapping(self, first, *filters)
+        assert len(states_) == len(filters) + 1, f'Expected {len(filters) + 1} states, got {len(states_)}'
+        if len(states_) == 1:
+            states = states_[0]
+        else:
+            states = tuple(states_)
+        return states  # type: ignore[bad-return-type]
+
+    @staticmethod
+    def merge(
+        state: Union['NestedDict[K, V]', 'FlattedDict[K, V]'],
+        /,
+        *states: Union['NestedDict[K, V]', 'FlattedDict[K, V]']
+    ) -> 'NestedDict[K, V]':
+        """
+        The inverse of :meth:`split()`.
+
+        ``merge`` takes one or more :class:`PrettyDict`'s and creates a new :class:`PrettyDict`.
+
+        Args:
+          state: A :class:`PrettyDict` object.
+          *states: Additional :class:`PrettyDict` objects.
+
+        Returns:
+          The merged :class:`PrettyDict`.
+        """
+        if not states:
+            return state
+        states = (state, *states)
+        new_state: FlattedDict = FlattedDict()
+        for state in states:
+            if isinstance(state, NestedDict):
+                new_state.update(state.to_flat())  # type: ignore[attribute-error] # pytype is wrong here
+            elif isinstance(state, FlattedDict):
+                new_state.update(state)
+            else:
+                raise TypeError(f'Expected Nested or Flatted Mapping, got {type(state)} instead.')
+        return NestedDict.from_flat(new_state)
+
+    def to_pure_dict(self) -> Dict[str, Any]:
+        flat_values = {k: x for k, x in self.to_flat().items()}
+        return nest_mapping(flat_values).to_dict()
+
+    def replace_by_pure_dict(
+        self,
+        pure_dict: Dict[str, Any],
+        replace_fn: Optional[SetValueFn] = None
+    ):
+        if replace_fn is None:
+            replace_fn = lambda x, v: x.replace(v) if hasattr(x, 'replace') else v
+        current_flat = self.to_flat()
+        for kp, v in flat_mapping(pure_dict).items():
+            if kp not in current_flat:
+                raise ValueError(f'key in pure_dict not available in state: {kp}')
+            current_flat[kp] = replace_fn(current_flat[kp], v)
+        self.update(nest_mapping(current_flat))
+
+
+class FlattedDict(PrettyDict):
+    """
+    A pytree-like structure that contains a :class:`Mapping` from strings or integers to leaves.
+
+    A valid leaf type is either :class:`State`, ``jax.Array``, ``numpy.ndarray`` or Python variables.
+
+    A :class:`NestedDict` can be generated by either calling :func:`states()` or
+    :func:`nodes()` on the :class:`Module`.
+
+    Example usage::
+
+      >>> import brainstate as brainstate
+      >>> import jax.numpy as jnp
+      >>>
+      >>> class Model(brainstate.nn.Module):
+      ...   def __init__(self):
+      ...     super().__init__()
+      ...     self.batchnorm = brainstate.nn.BatchNorm1d([10, 3])
+      ...     self.linear = brainstate.nn.Linear(2, 3)
+      ...   def __call__(self, x):
+      ...     return self.linear(self.batchnorm(x))
+      >>>
+      >>> model = Model()
+
+      >>> # retrieve the states of the model
+      >>> model.states()  # with the same to the function of ``brainstate.graph.states()``
+      FlattedDict({
+        ('batchnorm', 'running_mean'): LongTermState(
+          value=Array([[0., 0., 0.]], dtype=float32)
+        ),
+        ('batchnorm', 'running_var'): LongTermState(
+          value=Array([[1., 1., 1.]], dtype=float32)
+        ),
+        ('batchnorm', 'weight'): ParamState(
+          value={'bias': Array([[0., 0., 0.]], dtype=float32), 'scale': Array([[1., 1., 1.]], dtype=float32)}
+        ),
+        ('linear', 'weight'): ParamState(
+          value={'weight': Array([[-0.21467684,  0.7621282 , -0.50756454, -0.49047297],
+                 [-0.90413696,  0.6711    , -0.1254792 ,  0.50412565],
+                 [ 0.23975602,  0.47905368,  1.4851435 ,  0.16745673]],      dtype=float32), 'bias': Array([0., 0., 0., 0.], dtype=float32)}
+        )
+      })
+
+      >>> # retrieve the nodes of the model
+      >>> model.nodes()  # with the same to the function of ``brainstate.graph.nodes()``
+      FlattedDict({
+        ('batchnorm',): BatchNorm1d(
+          in_size=(10, 3),
+          out_size=(10, 3),
+          affine=True,
+          bias_initializer=Constant(value=0.0, dtype=<class 'numpy.float32'>),
+          scale_initializer=Constant(value=1.0, dtype=<class 'numpy.float32'>),
+          dtype=<class 'numpy.float32'>,
+          track_running_stats=True,
+          momentum=Array(shape=(), dtype=float32),
+          epsilon=Array(shape=(), dtype=float32),
+          feature_axis=(1,),
+          axis_name=None,
+          axis_index_groups=None,
+          running_mean=LongTermState(
+            value=Array(shape=(1, 3), dtype=float32)
+          ),
+          running_var=LongTermState(
+            value=Array(shape=(1, 3), dtype=float32)
+          ),
+          weight=ParamState(
+            value={'bias': Array(shape=(1, 3), dtype=float32), 'scale': Array(shape=(1, 3), dtype=float32)}
+          )
+        ),
+        ('linear',): Linear(
+          in_size=(10, 3),
+          out_size=(10, 4),
+          w_mask=None,
+          weight=ParamState(
+            value={'bias': Array(shape=(4,), dtype=float32), 'weight': Array(shape=(3, 4), dtype=float32)}
+          )
+        ),
+        (): Model(
+          batchnorm=BatchNorm1d(...),
+          linear=Linear(...)
+        )
+      })
+    """
+    __module__ = 'brainstate.util'
+
+    def __or__(self, other: 'FlattedDict[K, V]') -> 'FlattedDict[K, V]':
+        if not other:
+            return self
+        assert isinstance(other, FlattedDict), f'expected NestedDict; got {type(other).__qualname__}'
+        return FlattedDict.merge(self, other)
+
+    def __sub__(self, other: 'FlattedDict[K, V]') -> 'FlattedDict[K, V]':
+        if not other:
+            return self
+        assert isinstance(other, FlattedDict), f'expected NestedDict; got {type(other).__qualname__}'
+        diff = {k: v for k, v in self.items() if k not in other}
+        return FlattedDict(diff)
+
+    def to_nest(self) -> NestedDict:
+        """
+        Unflatten the flat mapping into a nested mapping.
+
+        Returns:
+            The nested mapping.
+        """
+        return nest_mapping(self)
+
+    @classmethod
+    def from_nest(
+        cls, nested_dict: abc.Mapping[PathParts, V] | Iterable[tuple[PathParts, V]],
+    ) -> 'FlattedDict':
+        """
+        Create a :class:`NestedDict` from a flat mapping.
+
+        Args:
+          nested_dict: The flat mapping.
+
+        Returns:
+          The :class:`NestedDict`.
+        """
+        return flat_mapping(nested_dict)
+
+    def split(  # type: ignore[misc]
+        self,
+        first: Filter,
+        /,
+        *filters: Filter
+    ) -> Union['FlattedDict[K, V]', tuple['FlattedDict[K, V]', ...]]:
+        """
+        Split a :class:`FlattedDict` into one or more :class:`FlattedDict`'s. The
+        user must pass at least one `:class:`Filter` (i.e. :class:`State`),
+        and the filters must be exhaustive (i.e. they must cover all
+        :class:`State` types in the :class:`NestedDict`).
+
+        Arguments:
+            first: The first filter
+            *filters: The optional, additional filters to group the state into mutually exclusive substates.
+
+        Returns:
+            One or more ``States`` equal to the number of filters passed.
+        """
+        filters = (first, *filters)
+        *states_, rest = _split_flatted_mapping(self, *filters)
+        if rest:
+            raise ValueError(f'Non-exhaustive filters, got a non-empty remainder: {rest}.\n'
+                             f'Use `...` to match all remaining elements.')
+
+        states: FlattedDict | Tuple[FlattedDict, ...]
+        if len(states_) == 1:
+            states = states_[0]
+        else:
+            states = tuple(states_)
+        return states  # type: ignore[bad-return-type]
+
+    def filter(
+        self,
+        first: Filter,
+        /,
+        *filters: Filter,
+    ) -> Union['FlattedDict[K, V]', Tuple['FlattedDict[K, V]', ...]]:
+        """
+        Filter a :class:`FlattedDict` into one or more :class:`FlattedDict`'s. The
+        user must pass at least one `:class:`Filter` (i.e. :class:`State`).
+        This method is similar to :meth:`split() <flax.nnx.NestedDict.state.split>`,
+        except the filters can be non-exhaustive.
+
+        Arguments:
+            first: The first filter
+            *filters: The optional, additional filters to group the state into mutually exclusive substates.
+
+        Returns:
+            One or more ``States`` equal to the number of filters passed.
+        """
+        *states_, _rest = _split_flatted_mapping(self, first, *filters)
+        assert len(states_) == len(filters) + 1, f'Expected {len(filters) + 1} states, got {len(states_)}'
+        if len(states_) == 1:
+            states = states_[0]
+        else:
+            states = tuple(states_)
+        return states  # type: ignore[bad-return-type]
+
+    @staticmethod
+    def merge(
+        state: Union['FlattedDict[K, V]', 'NestedDict[K, V]'],
+        /,
+        *states: Union['FlattedDict[K, V]', 'NestedDict[K, V]']
+    ) -> 'FlattedDict[K, V]':
+        """
+        The inverse of :meth:`split()`.
+
+        ``merge`` takes one or more :class:`FlattedDict`'s and creates a new :class:`FlattedDict`.
+
+        Args:
+          state: A :class:`PrettyDict` object.
+          *states: Additional :class:`PrettyDict` objects.
+
+        Returns:
+          The merged :class:`PrettyDict`.
+        """
+        if not states:
+            return state
+        states = (state, *states)
+        new_state: FlattedStateMapping[V] = {}
+        for state in states:
+            if isinstance(state, NestedDict):
+                new_state.update(state.to_flat())  # type: ignore[attribute-error] # pytype is wrong here
+            elif isinstance(state, FlattedDict):
+                new_state.update(state)
+            else:
+                raise TypeError(f'Expected Nested or Flatted Mapping, got {type(state)} instead.')
+        return FlattedDict(new_state)
+
+    def to_dict_values(self):
+        """
+        Convert a FlattedDict containing State objects to a plain dictionary of values.
+        
+        This method extracts the underlying values from any State objects in the FlattedDict,
+        creating a new dictionary with the same keys but where each State object is replaced
+        by its value attribute. Non-State objects are kept as is.
+        
+        Returns:
+            dict: A dictionary with the same keys as the FlattedDict, but where each State
+                  object is replaced by its value attribute. Non-State objects remain unchanged.
+        
+        Example:
+            >>> flat_dict = FlattedDict({('model', 'layer1', 'weight'): ParamState(value=jnp.ones((10, 5)))})
+            >>> flat_dict.to_dict_values()
+            {('model', 'layer1', 'weight'): Array([[1., 1., ...]], dtype=float32)}
+        """
+        from brainstate._state import State
+        return {
+            k: v.value if isinstance(v, State) else v
+            for k, v in self.items()
+        }
+
+    def assign_dict_values(self, data: dict):
+        """
+        Assign values from a dictionary to this FlattedDict.
+        
+        This method updates the values in the FlattedDict with values from the provided
+        dictionary. For keys that correspond to State objects, the value attribute of
+        the State is updated. For other keys, the value in the FlattedDict is directly
+        replaced with the new value.
+        
+        The method requires that all keys in the FlattedDict exist in the provided
+        dictionary, otherwise a KeyError is raised.
+        
+        Args:
+            data (dict): A dictionary containing the values to assign, where keys 
+                         must match those in the FlattedDict.
+        
+        Raises:
+            KeyError: If a key in the FlattedDict is not present in the provided dictionary.
+        
+        Example:
+            >>> flat_dict = FlattedDict({('model', 'weight'): ParamState(value=jnp.zeros((5, 5)))})
+            >>> flat_dict.assign_dict_values({('model', 'weight'): jnp.ones((5, 5))})
+            # The ParamState's value is now an array of ones
+        """
+        from brainstate._state import State
+        for k in self.keys():
+            if k not in data:
+                raise KeyError(f'Invalid key: {k!r}')
+            val = self[k]
+            if isinstance(val, State):
+                val.value = data[k]
+            else:
+                self[k] = data[k]
+
+
+def _split_nested_mapping(
+    mapping: 'NestedDict[K, V]',
+    *filters: Filter,
+) -> Tuple['NestedDict[K, V]', ...]:
+    # check if the filters are exhaustive
+    for i, filter_ in enumerate(filters):
+        if filter_ in (..., True) and i != len(filters) - 1:
+            remaining_filters = filters[i + 1:]
+            if not all(f in (..., True) for f in remaining_filters):
+                raise ValueError('`...` or `True` can only be used as the last filters, '
+                                 f'got {filter_} it at index {i}.')
+
+    # change the filters to predicates
+    predicates = tuple(map(to_predicate, filters))
+
+    # we have n + 1 state mappings, where n is the number of predicates
+    # the last state mapping is for values that don't match any predicate
+    flat_states: tuple[FlattedStateMapping[V], ...] = tuple({} for _ in range(len(predicates) + 1))
+
+    assert isinstance(mapping, NestedDict), f'expected NestedDict; got {type(mapping).__qualname__}'
+    flat_state = mapping.to_flat()
+    for path, value in flat_state.items():
+        for i, predicate in enumerate(predicates):
+            if predicate(path, value):
+                flat_states[i][path] = value  # type: ignore[index] # mypy is wrong here?
+                break
+        else:
+            # if we didn't break, set leaf to last state
+            flat_states[-1][path] = value  # type: ignore[index] # mypy is wrong here?
+
+    return tuple(NestedDict.from_flat(flat_state) for flat_state in flat_states)
+
+
+def _split_flatted_mapping(
+    mapping: FlattedDict[K, V],
+    *filters: Filter,
+) -> Tuple[FlattedDict[K, V], ...]:
+    # check if the filters are exhaustive
+    for i, filter_ in enumerate(filters):
+        if filter_ in (..., True) and i != len(filters) - 1:
+            remaining_filters = filters[i + 1:]
+            if not all(f in (..., True) for f in remaining_filters):
+                raise ValueError('`...` or `True` can only be used as the last filters, '
+                                 f'got {filter_} it at index {i}.')
+
+    # change the filters to predicates
+    predicates = tuple(map(to_predicate, filters))
+
+    # we have n + 1 state mappings, where n is the number of predicates
+    # the last state mapping is for values that don't match any predicate
+    flat_states: tuple[FlattedStateMapping[V], ...] = tuple({} for _ in range(len(predicates) + 1))
+
+    assert isinstance(mapping, FlattedDict), f'expected FlattedDict; got {type(mapping).__qualname__}'
+    for path, value in mapping.items():
+        for i, predicate in enumerate(predicates):
+            if predicate(path, value):
+                flat_states[i][path] = value  # type: ignore[index] # mypy is wrong here?
+                break
+        else:
+            # if we didn't break, set leaf to last state
+            flat_states[-1][path] = value  # type: ignore[index] # mypy is wrong here?
+
+    return tuple(FlattedDict(flat_state) for flat_state in flat_states)
+
+
+# register :class:`NestedDict` as a pytree
+def _nest_flatten_with_keys(x: NestedDict):
+    items = sorted(x.items())
+    children = tuple((jax.tree_util.DictKey(key), value) for key, value in items)
+    return children, tuple(key for key, _ in items)
+
+
+def _nest_unflatten(
+    static: Tuple[K, ...],
+    leaves: Tuple[V, ...] | Tuple[Dict[K, V]],
+):
+    return NestedDict(zip(static, leaves))
+
+
+jax.tree_util.register_pytree_with_keys(NestedDict,
+                                        _nest_flatten_with_keys,
+                                        _nest_unflatten)  # type: ignore[arg-type]
+
+
+# register :class:`FlattedDict` as a pytree
+
+def _flat_unflatten(
+    static: Tuple[K, ...],
+    leaves: Tuple[V, ...] | Tuple[Dict[K, V]],
+):
+    return FlattedDict(zip(static, leaves))
+
+
+jax.tree_util.register_pytree_with_keys(FlattedDict,
+                                        _nest_flatten_with_keys,
+                                        _flat_unflatten)  # type: ignore[arg-type]
+
+
+@jax.tree_util.register_pytree_node_class
+class PrettyList(list, PrettyRepr):
+    __module__ = 'brainstate.util'
+
+    def __pretty_repr__(self):
+        yield from yield_unique_pretty_repr_items(self, _list_repr_object, _list_repr_attr)
+
+    def __repr__(self):
+        return pretty_repr_object(self)
+
+    def tree_flatten(self):
+        return list(self), ()
+
+    @classmethod
+    def tree_unflatten(cls, aux_data, children):
+        return cls(children)
+
+
+def _list_repr_attr(node: PrettyList):
+    for v in node:
+        if isinstance(v, list):
+            v = PrettyList(v)
+        if isinstance(v, dict):
+            v = PrettyDict(v)
+        if isinstance(v, PrettyDict):
+            v = NestedStateRepr(v)
+        yield PrettyAttr('', v)
+
+
+def _list_repr_object(node: PrettyDict):
+    yield PrettyType('', value_sep='', start='[', end=']')
+
+
+def _repr_object_general(node: PrettyDict):
+    """
+    Generate a general representation of a PrettyDict object.
+
+    This function is used to create a pretty representation of a PrettyDict
+    object, which includes the type of the object and its value separator.
+
+    Args:
+        node (PrettyDict): The PrettyDict object to be represented.
+
+    Yields:
+        PrettyType: A PrettyType object representing the type of the node,
+        with specified value separator, start, and end characters.
+    """
+    yield PrettyType(type(node), value_sep='=', start='(', end=')')
+
+
+def _repr_attribute_general(node):
+    """
+    Generate a pretty representation of the attributes of a node.
+
+    This function iterates over the attributes of a given node and attempts
+    to generate a pretty representation for each attribute. It handles
+    conversion of lists and dictionaries to their pretty representation
+    counterparts and yields a PrettyAttr object for each attribute.
+
+    Args:
+        node: The object whose attributes are to be represented.
+
+    Yields:
+        PrettyAttr: A PrettyAttr object representing the key and value of
+        each attribute in a pretty format.
+    """
+    for k, v in vars(node).items():
+        try:
+            res = node.__pretty_repr_item__(k, v)
+            if res is None:
+                continue
+            k, v = res
+        except AttributeError:
+            pass
+
+        if k is None:
+            continue
+
+        # convert list to PrettyList
+        if isinstance(v, list):
+            v = PrettyList(v)
+
+        # convert dict to PrettyDict
+        if isinstance(v, dict):
+            v = PrettyDict(v)
+
+        # convert PrettyDict to NestedStateRepr
+        if isinstance(v, PrettyDict):
+            v = NestedStateRepr(v)
+
+        yield PrettyAttr(k, v)