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

brainstate/util/_pretty_pytree.py

@@ -0,0 +1,1301 @@
+# 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.
+
+"""Pretty printing utilities for PyTree-like structures.
+
+This module provides classes and functions for pretty printing and manipulating
+tree-like data structures in BrainState. The main components are:
+
+Classes:
+    - :class:`PrettyObject`: Base class for objects with pretty representation.
+    - :class:`PrettyDict`: Dictionary with pretty printing and tree utilities.
+    - :class:`NestedDict`: Nested mapping structure for hierarchical data.
+    - :class:`FlattedDict`: Flattened mapping with tuple keys for paths.
+    - :class:`PrettyList`: List with pretty printing capabilities.
+
+Functions:
+    - :func:`flat_mapping`: Flatten a nested mapping to tuple keys.
+    - :func:`nest_mapping`: Unflatten a mapping back to nested structure.
+
+All dictionary classes are registered as JAX pytrees and can be used with JAX
+transformations. They support splitting, filtering, and merging operations for
+organizing state in neural network models.
+
+Example:
+    >>> from brainstate.util import NestedDict, flat_mapping
+    >>> state = NestedDict({'layer1': {'weight': 1, 'bias': 2}})
+    >>> flat = state.to_flat()
+    >>> print(flat)
+    FlattedDict({('layer1', 'weight'): 1, ('layer1', 'bias'): 2})
+"""
+
+from collections import abc
+from typing import TypeVar, Hashable, Union, Iterable, Any, Optional, Tuple, Dict, Callable, Generator
+
+import jax
+
+from brainstate.typing import Filter, PathParts
+from ._pretty_repr import PrettyRepr, PrettyType, PrettyAttr, yield_unique_pretty_repr_items, pretty_repr_object
+from .filter import to_predicate
+from .struct import dataclass
+
+__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 = Callable[[Any], Any]
+SetValueFn = Callable[[V, Any], V]
+
+
+class PrettyObject(PrettyRepr):
+    """Base class for generating pretty representations of tree-like structures.
+
+    This class extends :class:`PrettyRepr` to provide a mechanism for generating
+    human-readable, pretty representations of tree-like data structures. It utilizes
+    custom functions to represent the object and its attributes in a structured and
+    visually appealing format.
+
+    The pretty representation is generated through the ``__pretty_repr__`` method,
+    which yields a sequence of pretty representation items using the general-purpose
+    representation functions.
+
+    Example:
+        >>> class MyTree(PrettyObject):
+        ...     def __init__(self, value):
+        ...         self.value = value
+        >>> tree = MyTree(42)
+        >>> print(tree)  # Uses __pretty_repr__ for display
+    """
+
+    def __pretty_repr__(self) -> Generator[Union[PrettyType, PrettyAttr], None, None]:
+        """Generate a pretty representation of the object.
+
+        This method yields a sequence of pretty representation items for the object,
+        using specialized functions to represent the object and its attributes.
+
+        Yields:
+            Union[PrettyType, PrettyAttr]: 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: Any, v: Any) -> Optional[Tuple[Any, Any]]:
+        """Transform a key-value pair 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. Subclasses
+        can override this method to customize how individual items are displayed.
+
+        Args:
+            k: The key of the item.
+            v: The value of the item.
+
+        Returns:
+            Optional[Tuple[Any, Any]]: A tuple containing the (key, value), or None
+                to skip this item in the representation.
+        """
+        return k, v
+
+
+PrettyReprTree = PrettyObject
+
+
+# The empty node is a struct.dataclass to be compatible with JAX.
+@dataclass
+class _EmptyNode:
+    """Sentinel class representing an empty node in tree structures."""
+    pass
+
+
+IsLeafCallable = Callable[[Tuple[Any, ...], abc.Mapping[Any, Any]], bool]
+_default_leaf: IsLeafCallable = lambda *args: False
+empty_node = _EmptyNode()
+
+
+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 into a flat mapping with tuple or string keys.
+
+    The nested keys are flattened to a tuple path. For example, ``{'a': {'b': 1}}``
+    becomes ``{('a', 'b'): 1}``. See :func:`nest_mapping` on how to restore the
+    nested structure.
+
+    Args:
+        xs: A nested mapping to flatten.
+        keep_empty_nodes: If True, replaces empty mappings with ``empty_node`` sentinel.
+            Otherwise, empty mappings are omitted from the result.
+        is_leaf: Optional function that takes ``(path, mapping)`` and returns True if
+            the mapping should be treated as a leaf (i.e., not flattened further).
+            Defaults to treating all mappings as non-leaves.
+        sep: If specified, keys in the returned mapping will be ``sep``-joined strings
+            instead of tuples. For example, with ``sep='/'``, ``('a', 'b')`` becomes ``'a/b'``.
+
+    Returns:
+        FlattedDict: A flattened mapping where nested keys are converted to tuples or strings.
+
+    Example:
+        >>> xs = {'foo': 1, 'bar': {'a': 2, 'b': {}}}
+        >>> flat_xs = flat_mapping(xs)
+        >>> flat_xs
+        FlattedDict({('foo',): 1, ('bar', 'a'): 2})
+
+        >>> # With separator
+        >>> flat_mapping(xs, sep='/')
+        FlattedDict({'foo': 1, 'bar/a': 2})
+
+        >>> # Keep empty nodes
+        >>> flat_mapping(xs, keep_empty_nodes=True)
+        FlattedDict({('foo',): 1, ('bar', 'a'): 2, ('bar', 'b'): _EmptyNode()})
+
+    Note:
+        Empty mappings are ignored by default and will not be restored by
+        :func:`nest_mapping` unless ``keep_empty_nodes=True``.
+    """
+    assert isinstance(xs, abc.Mapping), f'expected Mapping; got {type(xs).__qualname__}'
+
+    if sep is None:
+        def _key(path: Tuple[Any, ...]) -> Union[Tuple[Any, ...], str]:
+            return path
+    else:
+        def _key(path: Tuple[Any, ...]) -> Union[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: Optional[str] = None
+) -> 'NestedDict':
+    """Unflatten a mapping by converting tuple/string keys back to nested structure.
+
+    This is the inverse operation of :func:`flat_mapping`. It reconstructs a nested
+    mapping from a flattened one by interpreting tuple keys as paths in the nested
+    structure.
+
+    Args:
+        xs: A flattened mapping with tuple or string keys.
+        sep: Separator used to split string keys into paths. Must match the separator
+            used in :func:`flat_mapping()`. If None, keys are assumed to be tuples.
+
+    Returns:
+        NestedDict: A nested mapping reconstructed from the flattened structure.
+
+    Example:
+        >>> flat_xs = {
+        ...     ('foo',): 1,
+        ...     ('bar', 'a'): 2,
+        ... }
+        >>> xs = nest_mapping(flat_xs)
+        >>> xs
+        NestedDict({'foo': 1, 'bar': {'a': 2}})
+
+        >>> # With separator
+        >>> flat_xs_str = {'foo': 1, 'bar/a': 2}
+        >>> nest_mapping(flat_xs_str, sep='/')
+        NestedDict({'foo': 1, 'bar': {'a': 2}})
+
+    See Also:
+        :func:`flat_mapping`: The inverse operation that flattens a 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: Any, values: set) -> bool:
+    """Check if an object's identity is in a set of values.
+
+    Args:
+        x: The object to check.
+        values: A set of object identities to compare against.
+
+    Returns:
+        bool: True if the object's id is in the values set.
+    """
+    return id(x) in values
+
+
+def _default_process(x: Any) -> int:
+    """Get the identity of an object.
+
+    Args:
+        x: The object to process.
+
+    Returns:
+        int: The object's identity (id).
+    """
+    return id(x)
+
+
+class PrettyDict(dict, PrettyRepr):
+    """Base dictionary class with pretty representation and tree utilities.
+
+    This class extends the built-in dict with pretty printing capabilities and
+    provides base methods for tree operations. It serves as the parent class for
+    :class:`NestedDict` and :class:`FlattedDict`.
+
+    Attributes:
+        __module__ (str): Module identifier set to 'brainstate.util'.
+    """
+    __module__ = 'brainstate.util'
+
+    def __getattr__(self, key: K):  # type: ignore[misc]
+        """Access dictionary items as attributes.
+
+        Args:
+            key: The dictionary key to access.
+
+        Returns:
+            The value associated with the key.
+
+        Raises:
+            KeyError: If the key is not found in the dictionary.
+        """
+        return self[key]
+
+    def treefy_state(self) -> Any:
+        """Convert :class:`State` objects to a reference tree of the state.
+
+        This method traverses the tree structure and converts any :class:`State` objects
+        to their reference form using ``to_state_ref()``.
+
+        Returns:
+            Any: A tree structure where State objects are replaced with their references.
+
+        Example:
+            >>> from brainstate._state import State
+            >>> d = PrettyDict({'a': State(1), 'b': 2})
+            >>> ref_tree = d.treefy_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, Union[Dict[K, Any], V]]:
+        """Convert the :class:`PrettyDict` to a standard Python dictionary.
+
+        Returns:
+            Dict[K, Union[Dict[K, Any], V]]: A standard dictionary representation.
+        """
+        return dict(self)  # type: ignore
+
+    def __repr__(self) -> str:
+        """Generate a pretty string representation of the dictionary.
+
+        Returns:
+            str: A formatted string representation using pretty printing.
+        """
+        return pretty_repr_object(self)
+
+    def __pretty_repr__(self) -> Generator[Union[PrettyType, PrettyAttr], None, None]:
+        """Generate pretty representation items for this dictionary.
+
+        Yields:
+            Union[PrettyType, PrettyAttr]: Pretty representation items.
+        """
+        yield from yield_unique_pretty_repr_items(self, _default_repr_object, _default_repr_attr)
+
+    def split(self, *filters: Filter) -> Union['PrettyDict', Tuple['PrettyDict', ...]]:
+        """Split the dictionary based on filters (abstract method).
+
+        This method must be implemented by subclasses.
+
+        Args:
+            *filters: Filter specifications to split the dictionary.
+
+        Raises:
+            NotImplementedError: This is an abstract method.
+        """
+        raise NotImplementedError
+
+    def filter(self, *filters: Filter) -> Union['PrettyDict', Tuple['PrettyDict', ...]]:
+        """Filter the dictionary based on filters (abstract method).
+
+        This method must be implemented by subclasses.
+
+        Args:
+            *filters: Filter specifications to apply.
+
+        Raises:
+            NotImplementedError: This is an abstract method.
+        """
+        raise NotImplementedError
+
+    def merge(self, *states: 'PrettyDict') -> 'PrettyDict':
+        """Merge multiple dictionaries (abstract method).
+
+        This method must be implemented by subclasses.
+
+        Args:
+            *states: Additional PrettyDict objects to merge.
+
+        Raises:
+            NotImplementedError: This is an abstract method.
+        """
+        raise NotImplementedError
+
+    def subset(self, *filters: Filter) -> Union['PrettyDict', Tuple['PrettyDict', ...]]:
+        """Subset a :class:`PrettyDict` into one or more :class:`PrettyDict` instances.
+
+        The user must pass at least one :class:`Filter` (e.g., :class:`State`), and the
+        filters must be exhaustive (i.e., they must cover all :class:`State` types in
+        the :class:`PrettyDict`).
+
+        Args:
+            *filters: Filter specifications for subsetting.
+
+        Returns:
+            Union[PrettyDict, Tuple[PrettyDict, ...]]: One or more subsetted
+                dictionaries.
+        """
+        return self.filter(*filters)
+
+
+class NestedStateRepr(PrettyRepr):
+    """Pretty representation wrapper for nested state dictionaries.
+
+    This class wraps a :class:`PrettyDict` to provide specialized pretty printing
+    that displays the dictionary in a nested, compact format with curly braces.
+
+    Args:
+        state: The PrettyDict to wrap for pretty representation.
+
+    Attributes:
+        state (PrettyDict): The wrapped dictionary state.
+    """
+
+    def __init__(self, state: PrettyDict) -> None:
+        self.state = state
+
+    def __pretty_repr__(self) -> Generator[Union[PrettyType, PrettyAttr], None, None]:
+        """Generate a compact pretty representation of the nested state.
+
+        Yields:
+            Union[PrettyType, PrettyAttr]: Pretty representation items, skipping the
+                outer PrettyType from the wrapped state.
+        """
+        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: Any, subtree_renderer: Callable) -> Any:
+        """Generate a treescope representation for debugging tools.
+
+        Args:
+            path: The current path in the tree structure.
+            subtree_renderer: Callable to render subtrees.
+
+        Returns:
+            Any: Rendered representation of the nested state.
+        """
+        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) -> Generator[PrettyType, None, None]:
+    """Generate the default object representation for PrettyDict.
+
+    Args:
+        node: The PrettyDict node to represent.
+
+    Yields:
+        PrettyType: A type representation with dict-like formatting.
+    """
+    yield PrettyType('', value_sep=': ', start='{', end='}')
+
+
+def _default_repr_attr(node: PrettyDict) -> Generator[PrettyAttr, None, None]:
+    """Generate the default attribute representations for PrettyDict items.
+
+    This function converts list and dict values to their pretty equivalents
+    and wraps PrettyDict values in NestedStateRepr for compact display.
+
+    Args:
+        node: The PrettyDict node whose attributes to represent.
+
+    Yields:
+        PrettyAttr: Pretty attribute representations for each item.
+    """
+    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 nested mapping structure for organizing hierarchical data.
+
+    This class represents a nested mapping from strings or integers to leaves, where
+    valid leaf types include :class:`State`, ``jax.Array``, ``numpy.ndarray``, or
+    nested :class:`NestedDict` and :class:`FlattedDict` structures.
+
+    :class:`NestedDict` is a JAX pytree and can be used with JAX transformations.
+    It provides methods for flattening to :class:`FlattedDict`, splitting/filtering
+    based on predicates, and merging multiple nested structures.
+
+    Attributes:
+        __module__ (str): Module identifier set to 'brainstate.util'.
+
+    Example:
+        >>> from brainstate.util import NestedDict
+        >>> state = NestedDict({
+        ...     'layer1': {'weight': jnp.ones((3, 3)), 'bias': jnp.zeros(3)},
+        ...     'layer2': {'weight': jnp.ones((3, 1))}
+        ... })
+        >>> flat = state.to_flat()
+        >>> print(flat)
+        FlattedDict({('layer1', 'weight'): ..., ('layer1', 'bias'): ..., ...})
+
+    See Also:
+        :class:`FlattedDict`: The flattened counterpart with tuple keys.
+        :func:`flat_mapping`: Function to flatten a nested mapping.
+        :func:`nest_mapping`: Function to unflatten a flat mapping.
+    """
+    __module__ = 'brainstate.util'
+
+    def __or__(self, other: 'NestedDict') -> 'NestedDict':
+        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') -> 'NestedDict':
+        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(self, *filters: Filter) -> Union['NestedDict', Tuple['NestedDict', ...]]:
+        """
+        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.
+        """
+        *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, *filters: Filter) -> Union['NestedDict', Tuple['NestedDict', ...]]:
+        """
+        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_ = _split_nested_mapping(self, *filters)
+        assert len(states_) == len(filters) + 1, f'Expected {len(filters) + 1} states, got {len(states_)}'
+        if len(filters) == 1:
+            return states_[0]
+        else:
+            return tuple(states_[:-1])
+
+    @staticmethod
+    def merge(*states) -> 'NestedDict':
+        """
+        The inverse of :meth:`split()`.
+
+        ``merge`` takes one or more :class:`PrettyDict`'s and creates a new :class:`PrettyDict`.
+
+        Args:
+          *states: Additional :class:`PrettyDict` objects.
+
+        Returns:
+          The merged :class:`PrettyDict`.
+        """
+        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]:
+        """Convert to a pure nested dictionary structure.
+
+        This method creates a standard Python dictionary with the same nested structure
+        as this NestedDict, without any special class wrappers.
+
+        Returns:
+            Dict[str, Any]: A pure nested dictionary representation.
+
+        Example:
+            >>> nested = NestedDict({'a': {'b': 1, 'c': 2}})
+            >>> pure = nested.to_pure_dict()
+            >>> type(pure)
+            <class 'dict'>
+        """
+        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
+    ) -> None:
+        """Replace values in this NestedDict using a pure dictionary.
+
+        This method updates the values in this NestedDict with values from a standard
+        Python dictionary. For :class:`State` objects with a ``replace`` method, the
+        replace method is called; otherwise, values are directly assigned.
+
+        Args:
+            pure_dict: A pure dictionary with matching structure containing new values.
+            replace_fn: Optional custom function to replace values. Takes ``(old_value, new_value)``
+                and returns the updated value. Defaults to calling ``replace()`` method if
+                available, otherwise direct assignment.
+
+        Raises:
+            ValueError: If a key in ``pure_dict`` is not found in this NestedDict.
+
+        Example:
+            >>> from brainstate._state import State
+            >>> nested = NestedDict({'a': State(1), 'b': 2})
+            >>> nested.replace_by_pure_dict({'a': 10, 'b': 20})
+            >>> nested['a'].value
+            10
+        """
+        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') -> 'FlattedDict':
+        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') -> 'FlattedDict':
+        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', tuple['FlattedDict', ...]]:
+        """
+        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', Tuple['FlattedDict', ...]]:
+        """
+        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(*states: Union['FlattedDict', 'NestedDict']) -> 'FlattedDict':
+        """
+        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`.
+        """
+        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) -> Dict[PathParts, Any]:
+        """Convert a FlattedDict containing State objects to a dictionary of raw values.
+
+        This method extracts the underlying values from any :class:`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[PathParts, Any]: A dictionary with the same keys as the FlattedDict, but
+                where each State object is replaced by its value. Non-State objects remain
+                unchanged.
+
+        Example:
+            >>> from brainstate._state import ParamState
+            >>> flat_dict = FlattedDict({
+            ...     ('model', 'layer1', 'weight'): ParamState(value=jnp.ones((10, 5)))
+            ... })
+            >>> values = flat_dict.to_dict_values()
+            >>> 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[PathParts, Any]) -> None:
+        """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 :class:`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.
+
+        Args:
+            data: 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:
+            >>> from brainstate._state import ParamState
+            >>> 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',
+    *filters: Filter,
+) -> Tuple['NestedDict', ...]:
+    """Split a nested mapping into multiple nested mappings based on filters.
+
+    This internal function partitions a NestedDict into multiple NestedDicts based on
+    filter predicates. Items that match each filter are placed in separate mappings,
+    with unmatched items going to the final mapping.
+
+    Args:
+        mapping: The NestedDict to split.
+        *filters: Filter specifications. The catch-all filters ``...`` or ``True`` can
+            only be used as the last filter.
+
+    Returns:
+        Tuple[NestedDict, ...]: A tuple of n+1 NestedDicts, where n is the number
+            of filters. The last mapping contains items that didn't match any filter.
+
+    Raises:
+        ValueError: If ``...`` or ``True`` is used before the last filter position.
+        AssertionError: If mapping is not a NestedDict.
+    """
+    # 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]
+                break
+        else:
+            # If we didn't break, set leaf to last state
+            flat_states[-1][path] = value  # type: ignore[index]
+
+    return tuple(NestedDict.from_flat(flat_state) for flat_state in flat_states)
+
+
+def _split_flatted_mapping(
+    mapping: FlattedDict,
+    *filters: Filter,
+) -> Tuple[FlattedDict, ...]:
+    """Split a flattened mapping into multiple flattened mappings based on filters.
+
+    This internal function partitions a FlattedDict into multiple FlattedDicts based on
+    filter predicates. Items that match each filter are placed in separate mappings,
+    with unmatched items going to the final mapping.
+
+    Args:
+        mapping: The FlattedDict to split.
+        *filters: Filter specifications. The catch-all filters ``...`` or ``True`` can
+            only be used as the last filter.
+
+    Returns:
+        Tuple[FlattedDict, ...]: A tuple of n+1 FlattedDicts, where n is the number
+            of filters. The last mapping contains items that didn't match any filter.
+
+    Raises:
+        ValueError: If ``...`` or ``True`` is used before the last filter position.
+        AssertionError: If mapping is not a FlattedDict.
+    """
+    # 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]
+                break
+        else:
+            # If we didn't break, set leaf to last state
+            flat_states[-1][path] = value  # type: ignore[index]
+
+    return tuple(FlattedDict(flat_state) for flat_state in flat_states)
+
+
+# Register :class:`NestedDict` as a pytree
+def _nest_flatten_with_keys(x: NestedDict) -> Tuple[Tuple[Tuple[jax.tree_util.DictKey, Any], ...], Tuple[K, ...]]:
+    """Flatten a NestedDict for JAX pytree registration with keys.
+
+    Args:
+        x: The NestedDict to flatten.
+
+    Returns:
+        Tuple containing:
+
+            - Tuple of (key, value) pairs where keys are wrapped in DictKey
+            - Tuple of static keys for reconstruction
+    """
+    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: Union[Tuple[V, ...], Tuple[Dict]],
+) -> NestedDict:
+    """Unflatten a NestedDict from pytree components.
+
+    Args:
+        static: Tuple of keys for reconstruction.
+        leaves: Tuple of leaf values.
+
+    Returns:
+        NestedDict: Reconstructed NestedDict.
+    """
+    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: Union[Tuple[V, ...], Tuple[Dict]],
+) -> FlattedDict:
+    """Unflatten a FlattedDict from pytree components.
+
+    Args:
+        static: Tuple of keys for reconstruction.
+        leaves: Tuple of leaf values.
+
+    Returns:
+        FlattedDict: Reconstructed FlattedDict.
+    """
+    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):
+    """A list subclass with pretty representation and JAX pytree compatibility.
+
+    This class extends the built-in list with pretty printing capabilities and
+    registers itself as a JAX pytree for use in JAX transformations.
+
+    Attributes:
+        __module__ (str): Module identifier set to 'brainstate.util'.
+
+    Example:
+        >>> from brainstate.util import PrettyList
+        >>> lst = PrettyList([1, 2, {'a': 3}])
+        >>> print(lst)  # Pretty formatted output
+        [1, 2, {'a': 3}]
+    """
+    __module__ = 'brainstate.util'
+
+    def __pretty_repr__(self) -> Generator[Union[PrettyType, PrettyAttr], None, None]:
+        """Generate pretty representation items for this list.
+
+        Yields:
+            Union[PrettyType, PrettyAttr]: Pretty representation items.
+        """
+        yield from yield_unique_pretty_repr_items(self, _list_repr_object, _list_repr_attr)
+
+    def __repr__(self) -> str:
+        """Generate a pretty string representation of the list.
+
+        Returns:
+            str: A formatted string representation using pretty printing.
+        """
+        return pretty_repr_object(self)
+
+    def tree_flatten(self) -> Tuple[list, Tuple]:
+        """Flatten the list for JAX pytree operations.
+
+        Returns:
+            Tuple containing:
+                - The list items as children
+                - Empty tuple as auxiliary data
+        """
+        return list(self), ()
+
+    @classmethod
+    def tree_unflatten(cls, aux_data: Tuple, children: list) -> 'PrettyList':
+        """Reconstruct a PrettyList from pytree components.
+
+        Args:
+            aux_data: Auxiliary data (unused).
+            children: List items to reconstruct from.
+
+        Returns:
+            PrettyList: Reconstructed PrettyList.
+        """
+        return cls(children)
+
+
+def _list_repr_attr(node: PrettyList) -> Generator[PrettyAttr, None, None]:
+    """Generate attribute representations for PrettyList items.
+
+    This function converts list and dict values to their pretty equivalents
+    and wraps PrettyDict values in NestedStateRepr for compact display.
+
+    Args:
+        node: The PrettyList whose items to represent.
+
+    Yields:
+        PrettyAttr: Pretty attribute representations for each item with empty keys.
+    """
+    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: PrettyList) -> Generator[PrettyType, None, None]:
+    """Generate the object representation for PrettyList.
+
+    Args:
+        node: The PrettyList to represent.
+
+    Yields:
+        PrettyType: A type representation with list-like formatting.
+    """
+    yield PrettyType('', value_sep='', start='[', end=']')
+
+
+def _repr_object_general(node: Any) -> Generator[PrettyType, None, None]:
+    """Generate a general object representation for any PrettyObject.
+
+    This function creates a pretty representation of an object that includes
+    the type of the object with parentheses formatting.
+
+    Args:
+        node: The 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: Any) -> Generator[PrettyAttr, None, None]:
+    """Generate a pretty representation of the attributes of a general object.
+
+    This function iterates over the attributes of a given node and generates
+    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 visible attribute.
+
+    The function respects the ``__pretty_repr_item__`` method if available on
+    the node, allowing custom filtering or transformation of attribute items.
+
+    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)