brainstate 0.1.0.post20250210__py2.py3-none-any.whl → 0.1.0.post20250212__py2.py3-none-any.whl

brainstate/_state.py

@@ -19,8 +19,10 @@ import contextlib
 import dataclasses
 import threading
 from functools import wraps, partial
-from typing import (Any, Union, Callable, Generic, Mapping,
-                    TypeVar, Optional, TYPE_CHECKING, Tuple, Dict, List, Sequence)
+from typing import (
+    Any, Union, Callable, Generic, Mapping,
+    TypeVar, Optional, TYPE_CHECKING, Tuple, Dict, List, Sequence
+)
 
 import jax
 import numpy as np
@@ -28,7 +30,7 @@ from jax.api_util import shaped_abstractify
 from jax.extend import source_info_util
 
 from brainstate.typing import ArrayLike, PyTree, Missing
-from brainstate.util import DictManager, PrettyRepr, PrettyType, PrettyAttr
+from brainstate.util import DictManager, PrettyObject
 
 __all__ = [
     'State', 'ShortTermState', 'LongTermState', 'HiddenState', 'ParamState', 'TreefyState',
@@ -184,7 +186,7 @@ def _get_trace_stack_level() -> int:
     return len(TRACE_CONTEXT.state_stack)
 
 
-class State(Generic[A], PrettyRepr):
+class State(Generic[A], PrettyObject):
     """
     The pointer to specify the dynamical data.
 
@@ -434,45 +436,25 @@ class State(Generic[A], PrettyRepr):
         del metadata['_value']
         return TreefyState(type(self), self._value, **metadata)
 
-    def __pretty_repr__(self):
-        yield PrettyType(type=type(self))
-        for name, value in vars(self).items():
-            if name == '_value':
-                name = 'value'
-            if name == '_name':
-                if value is None:
-                    continue
-                else:
-                    name = 'name'
-            if name == 'tag' and value is None:
-                continue
-            if name in ['_level', '_source_info', '_been_writen']:
-                continue
-            yield PrettyAttr(name, repr(value))
-
-    def __treescope_repr__(self, path, subtree_renderer):
-        children = {}
-        for name, value in vars(self).items():
-            if name == '_value':
-                name = 'value'
-            if name == '_name':
-                if value is None:
-                    continue
-                else:
-                    name = 'name'
-            if name == 'tag' and value is None:
-                continue
-            if name in ['_level', '_source_info', '_been_writen']:
-                continue
-            children[name] = value
-
-        import treescope  # type: ignore[import-not-found,import-untyped]
-        return treescope.repr_lib.render_object_constructor(
-            object_type=type(self),
-            attributes=children,
-            path=path,
-            subtree_renderer=subtree_renderer,
-        )
+    def __pretty_repr_item__(self, k, v):
+        if k in ['_level', '_source_info', '_been_writen']:
+            return None, None
+        if k == '_value':
+            return 'value', v
+
+        if k == '_name':
+            if self.name is None:
+                return None, None
+            else:
+                return 'name', v
+
+        if k == 'tag':
+            if self.tag is None:
+                return None, None
+            else:
+                return 'tag', v
+
+        return k, v
 
     def __eq__(self, other: object) -> bool:
         return type(self) is type(other) and vars(other) == vars(self)
@@ -483,6 +465,25 @@ class State(Generic[A], PrettyRepr):
         """
         return hash(id(self))
 
+    def numel(self) -> int:
+        """
+        Calculate the total number of elements in the state value.
+
+        This method traverses the state's value, which may be a nested structure (PyTree),
+        and computes the sum of sizes of all leaf nodes.
+
+        Returns:
+            int: The total number of elements across all arrays in the state value.
+                For scalar values, this will be 1. For arrays or nested structures,
+                it will be the sum of the sizes of all contained arrays.
+
+        Note:
+            This method uses jax.tree.leaves to flatten any nested structure in the state value,
+            and jax.numpy.size to compute the size of each leaf node.
+        """
+        sizes = [jax.numpy.size(val) for val in jax.tree.leaves(self._value)]
+        return sum(sizes)
+
 
 def record_state_init(st: State[A]):
     trace: Catcher
@@ -809,7 +810,7 @@ class StateTraceStack(Generic[A]):
         return StateTraceStack().merge(self, other)
 
 
-class TreefyState(Generic[A], PrettyRepr):
+class TreefyState(Generic[A], PrettyObject):
     """
     The state as a pytree.
     """
@@ -831,35 +832,15 @@ class TreefyState(Generic[A], PrettyRepr):
 
         def __delattr__(self, name: str) -> None: ...
 
-    def __pretty_repr__(self):
-        yield PrettyType(type=type(self))
-        yield PrettyAttr('type', self.type.__name__)
-        for name, value in vars(self).items():
-            if name == '_value':
-                name = 'value'
-            if name == '_name':
-                if value is None:
-                    continue
-                else:
-                    name = 'name'
-            if name in ['_level', '_source_info', 'type']:
-                continue
-            yield PrettyAttr(name, repr(value))
-
-    def __treescope_repr__(self, path, subtree_renderer):
-        children = {'type': self.type}
-        for name, value in vars(self).items():
-            if name == 'type':
-                continue
-            children[name] = value
-
-        import treescope  # type: ignore[import-not-found,import-untyped]
-        return treescope.repr_lib.render_object_constructor(
-            object_type=type(self),
-            attributes=children,
-            path=path,
-            subtree_renderer=subtree_renderer,
-        )
+    def __pretty_repr_item__(self, k, v):
+        if k in ['_level', '_source_info', '_been_writen']:
+            return None, None
+        if k == '_value':
+            return 'value', v
+
+        if k == '_name':
+            return (None, None) if v is None else ('name', v)
+        return k, v
 
     def replace(self, value: B) -> TreefyState[B]:
         """

brainstate/compile/_make_jaxpr.py

@@ -206,7 +206,7 @@ class StatefulFunction(object):
         self._cached_state_trace: Dict[Any, StateTraceStack] = dict()
 
     def __repr__(self) -> str:
-        return (f"{self.__class__.__name__}({self.fun}, "
+        return (f"{self.__class__.__name__}("
                 f"static_argnums={self.static_argnums}, "
                 f"axis_env={self.axis_env}, "
                 f"abstracted_axes={self.abstracted_axes}, "

brainstate/graph/_graph_node.py

@@ -27,7 +27,7 @@ import numpy as np
 
 from brainstate._state import State, TreefyState
 from brainstate.typing import Key
-from brainstate.util._pretty_repr import PrettyRepr, pretty_repr_avoid_duplicate, PrettyType, PrettyAttr
+from brainstate.util._pretty_repr import PrettyRepr, yield_unique_pretty_repr_items, PrettyType, PrettyAttr
 from ._graph_operation import register_graph_node_type
 
 __all__ = [
@@ -88,7 +88,7 @@ class Node(PrettyRepr, metaclass=GraphNodeMeta):
         """
         Pretty repr for the object.
         """
-        yield from pretty_repr_avoid_duplicate(self, _default_repr_object, _default_repr_attr)
+        yield from yield_unique_pretty_repr_items(self, _default_repr_object, _default_repr_attr)
 
     def __treescope_repr__(self, path, subtree_renderer):
         """

brainstate/nn/_interaction/_conv.py

@@ -193,16 +193,18 @@ class _Conv(_BaseConv):
         name: str = None,
         param_type: type = ParamState,
     ):
-        super().__init__(in_size=in_size,
-                         out_channels=out_channels,
-                         kernel_size=kernel_size,
-                         stride=stride,
-                         padding=padding,
-                         lhs_dilation=lhs_dilation,
-                         rhs_dilation=rhs_dilation,
-                         groups=groups,
-                         w_mask=w_mask,
-                         name=name)
+        super().__init__(
+            in_size=in_size,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            lhs_dilation=lhs_dilation,
+            rhs_dilation=rhs_dilation,
+            groups=groups,
+            w_mask=w_mask,
+            name=name
+        )
 
         self.w_initializer = w_init
         self.b_initializer = b_init

brainstate/nn/_module.py

@@ -294,7 +294,8 @@ class Sequential(Module):
         # the input and output shape
         if first.in_size is not None:
             self.in_size = first.in_size
-        self.out_size = tuple(in_size)
+        if in_size is not None:
+            self.out_size = tuple(in_size)
 
     def update(self, x):
         """Update function of a sequential model.

brainstate/surrogate.py

@@ -21,6 +21,8 @@ import jax.numpy as jnp
 import jax.scipy as sci
 from jax.interpreters import batching, ad, mlir
 
+from brainstate.util._dict import PrettyObject
+
 if jax.__version_info__ < (0, 4, 38):
     from jax.core import Primitive
 else:
@@ -77,7 +79,7 @@ def _heaviside_imp(x, dx):
 
 
 def _heaviside_batching(args, axes):
-    return heaviside_p.bind(*args), [axes[0]]
+    return heaviside_p.bind(*args), tuple(axes)
 
 
 def _heaviside_jvp(primals, tangents):
@@ -97,7 +99,7 @@ ad.primitive_jvps[heaviside_p] = _heaviside_jvp
 mlir.register_lowering(heaviside_p, mlir.lower_fun(_heaviside_imp, multiple_results=True))
 
 
-class Surrogate(object):
+class Surrogate(PrettyObject):
     """The base surrograte gradient function.
 
     To customize a surrogate gradient function, you can inherit this class and
@@ -142,9 +144,20 @@ class Surrogate(object):
 class Sigmoid(Surrogate):
     """Spike function with the sigmoid-shaped surrogate gradient.
 
+    This class implements a spiking neuron activation with a sigmoid-shaped
+    surrogate gradient for backpropagation. It can be used in spiking neural
+    networks to approximate the non-differentiable step function during training.
+
+    Parameters
+    ----------
+    alpha : float, optional
+        A parameter controlling the steepness of the sigmoid curve in the
+        surrogate gradient. Higher values make the transition sharper.
+        Default is 4.0.
+
     See Also
     --------
-    sigmoid
+    sigmoid : Function version of this class.
 
     """
 
@@ -153,9 +166,33 @@ class Sigmoid(Surrogate):
         self.alpha = alpha
 
     def surrogate_fun(self, x):
+        """Compute the surrogate function.
+
+        Parameters
+        ----------
+        x : jax.Array
+            The input array.
+
+        Returns
+        -------
+        jax.Array
+            The output of the surrogate function.
+        """
         return sci.special.expit(self.alpha * x)
 
     def surrogate_grad(self, x):
+        """Compute the gradient of the surrogate function.
+
+        Parameters
+        ----------
+        x : jax.Array
+            The input array.
+
+        Returns
+        -------
+        jax.Array
+            The gradient of the surrogate function.
+        """
         sgax = sci.special.expit(x * self.alpha)
         dx = (1. - sgax) * sgax * self.alpha
         return dx
@@ -171,7 +208,12 @@ def sigmoid(
     x: jax.Array,
     alpha: float = 4.,
 ):
-    r"""Spike function with the sigmoid-shaped surrogate gradient.
+    r"""
+    Compute a spike function with a sigmoid-shaped surrogate gradient.
+
+    This function implements a spiking neuron activation with a sigmoid-shaped
+    surrogate gradient for backpropagation. It can be used in spiking neural
+    networks to approximate the non-differentiable step function during training.
 
     If `origin=False`, return the forward function:
 
@@ -210,16 +252,28 @@ def sigmoid(
 
     Parameters
     ----------
-    x: jax.Array, Array
-      The input data.
-    alpha: float
-      Parameter to control smoothness of gradient
-
+    x : jax.Array
+        The input array representing the neuron's membrane potential.
+    alpha : float, optional
+        A parameter controlling the steepness of the sigmoid curve in the
+        surrogate gradient. Higher values make the transition sharper.
+        Default is 4.0.
 
     Returns
     -------
-    out: jax.Array
-      The spiking state.
+    jax.Array
+        An array of the same shape as the input, containing binary values (0 or 1)
+        representing the spiking state of each neuron.
+        
+    Notes
+    -----
+    The forward pass uses a step function (1 for x >= 0, 0 for x < 0),
+    while the backward pass uses a sigmoid-shaped surrogate gradient for
+    smooth optimization.
+
+    The surrogate gradient is defined as:
+    g'(x) = alpha * (1 - sigmoid(alpha * x)) * sigmoid(alpha * x)
+
     """
     return Sigmoid(alpha=alpha)(x)
 
@@ -238,11 +292,15 @@ class PiecewiseQuadratic(Surrogate):
         self.alpha = alpha
 
     def surrogate_fun(self, x):
-        z = jnp.where(x < -1 / self.alpha,
-                      0.,
-                      jnp.where(x > 1 / self.alpha,
-                                1.,
-                                (-self.alpha * jnp.abs(x) / 2 + 1) * self.alpha * x + 0.5))
+        z = jnp.where(
+            x < -1 / self.alpha,
+            0.,
+            jnp.where(
+                x > 1 / self.alpha,
+                1.,
+                (-self.alpha * jnp.abs(x) / 2 + 1) * self.alpha * x + 0.5
+            )
+        )
         return z
 
     def surrogate_grad(self, x):
@@ -260,7 +318,12 @@ def piecewise_quadratic(
     x: jax.Array,
     alpha: float = 1.,
 ):
-    r"""Judge spiking state with a piecewise quadratic function [1]_ [2]_ [3]_ [4]_ [5]_.
+    r"""
+    Judge spiking state with a piecewise quadratic function [1]_ [2]_ [3]_ [4]_ [5]_.
+
+    This function implements a surrogate gradient method for spiking neural networks
+    using a piecewise quadratic function. It provides a differentiable approximation
+    of the step function used in the forward pass of spiking neurons.
 
     If `origin=False`, computes the forward function:
 
@@ -306,18 +369,29 @@ def piecewise_quadratic(
        >>> plt.legend()
        >>> plt.show()
 
-    Parameters
+  Parameters
     ----------
-    x: jax.Array, Array
-      The input data.
-    alpha: float
-      Parameter to control smoothness of gradient
-
+    x : jax.Array
+        The input array representing the neuron's membrane potential.
+    alpha : float, optional
+        A parameter controlling the steepness of the surrogate gradient.
+        Higher values result in a steeper gradient. Default is 1.0.
 
     Returns
     -------
-    out: jax.Array
-      The spiking state.
+    jax.Array
+        An array of the same shape as the input, containing binary values (0 or 1)
+        representing the spiking state of each neuron.
+
+    Notes
+    -----
+    The function uses different computations for forward and backward passes:
+    - Forward: Step function (1 for x >= 0, 0 for x < 0)
+    - Backward: Piecewise quadratic function for smooth gradient
+
+    The surrogate gradient is defined as:
+    g'(x) = 0 if |x| > 1/alpha
+            -alpha^2|x| + alpha if |x| <= 1/alpha
 
     References
     ----------
@@ -331,11 +405,22 @@ def piecewise_quadratic(
 
 
 class PiecewiseExp(Surrogate):
-    """Judge spiking state with a piecewise exponential function.
+    """
+    Judge spiking state with a piecewise exponential function.
+
+    This class implements a surrogate gradient method for spiking neural networks
+    using a piecewise exponential function. It provides a differentiable approximation
+    of the step function used in the forward pass of spiking neurons.
+
+    Parameters
+    ----------
+    alpha : float, optional
+        A parameter controlling the steepness of the surrogate gradient.
+        Higher values result in a steeper gradient. Default is 1.0.
 
     See Also
     --------
-    piecewise_exp
+    piecewise_exp : Function version of this class.
     """
 
     def __init__(self, alpha: float = 1.):
@@ -343,16 +428,62 @@ class PiecewiseExp(Surrogate):
         self.alpha = alpha
 
     def surrogate_grad(self, x):
+        """
+        Compute the surrogate gradient.
+
+        Parameters
+        ----------
+        x : jax.Array
+            The input array.
+
+        Returns
+        -------
+        jax.Array
+            The surrogate gradient.
+        """
         dx = (self.alpha / 2) * jnp.exp(-self.alpha * jnp.abs(x))
         return dx
 
     def surrogate_fun(self, x):
-        return jnp.where(x < 0, jnp.exp(self.alpha * x) / 2, 1 - jnp.exp(-self.alpha * x) / 2)
+        """
+        Compute the surrogate function.
+
+        Parameters
+        ----------
+        x : jax.Array
+            The input array.
+
+        Returns
+        -------
+        jax.Array
+            The output of the surrogate function.
+        """
+        return jnp.where(
+            x < 0,
+            jnp.exp(self.alpha * x) / 2,
+            1 - jnp.exp(-self.alpha * x) / 2
+        )
 
     def __repr__(self):
+        """
+        Return a string representation of the PiecewiseExp instance.
+
+        Returns
+        -------
+        str
+            A string representation of the instance.
+        """
         return f'{self.__class__.__name__}(alpha={self.alpha})'
 
     def __hash__(self):
+        """
+        Compute a hash value for the PiecewiseExp instance.
+
+        Returns
+        -------
+        int
+            A hash value for the instance.
+        """
         return hash((self.__class__, self.alpha))
 
 
@@ -363,6 +494,10 @@ def piecewise_exp(
 ):
     r"""Judge spiking state with a piecewise exponential function [1]_.
 
+    This function implements a surrogate gradient method for spiking neural networks
+    using a piecewise exponential function. It provides a differentiable approximation
+    of the step function used in the forward pass of spiking neurons.
+
     If `origin=False`, computes the forward function:
 
     .. math::
@@ -403,16 +538,26 @@ def piecewise_exp(
 
     Parameters
     ----------
-    x: jax.Array, Array
-      The input data.
-    alpha: float
-      Parameter to control smoothness of gradient
-
+    x : jax.Array
+        The input array representing the neuron's membrane potential.
+    alpha : float, optional
+        A parameter controlling the steepness of the surrogate gradient.
+        Higher values result in a steeper gradient. Default is 1.0.
 
     Returns
     -------
-    out: jax.Array
-      The spiking state.
+    jax.Array
+        An array of the same shape as the input, containing binary values (0 or 1)
+        representing the spiking state of each neuron.
+
+    Notes
+    -----
+    The function uses different computations for forward and backward passes:
+    - Forward: Step function (1 for x >= 0, 0 for x < 0)
+    - Backward: Piecewise exponential function for smooth gradient
+
+    The surrogate gradient is defined as:
+    g'(x) = (alpha / 2) * exp(-alpha * |x|)
 
     References
     ----------

brainstate/util/_dict.py

@@ -24,11 +24,17 @@ import jax
 
 from brainstate.typing import Filter, PathParts
 from ._filter import to_predicate
-from ._pretty_repr import PrettyRepr, PrettyType, PrettyAttr, pretty_repr_avoid_duplicate, get_repr
+from ._pretty_repr import PrettyRepr, PrettyType, PrettyAttr, yield_unique_pretty_repr_items, pretty_repr_object
 from ._struct import dataclass
 
 __all__ = [
-    'NestedDict', 'FlattedDict', 'flat_mapping', 'nest_mapping',
+    'PrettyDict',
+    'NestedDict',
+    'FlattedDict',
+    'flat_mapping',
+    'nest_mapping',
+    'PrettyList',
+    'PrettyObject',
 ]
 
 A = TypeVar('A')
@@ -40,6 +46,119 @@ ExtractValueFn = abc.Callable[[Any], Any]
 SetValueFn = abc.Callable[[V, Any], V]
 
 
+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)
+
+
+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:
@@ -213,10 +332,10 @@ class PrettyDict(dict, PrettyRepr):
 
     def __repr__(self) -> str:
         # repr the individual object with the pretty representation
-        return get_repr(self)
+        return pretty_repr_object(self)
 
     def __pretty_repr__(self):
-        yield from pretty_repr_avoid_duplicate(self, _default_repr_object, _default_repr_attr)
+        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
@@ -237,15 +356,20 @@ class PrettyDict(dict, PrettyRepr):
 
 
 def _default_repr_object(node: PrettyDict):
-    yield PrettyType(type(node), value_sep=': ', start='({', end='})')
+    yield PrettyType('', value_sep=': ', start='{', end='}')
 
 
-def _default_repr_attr(node: PrettyDict):
+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)
 
 
@@ -735,3 +859,36 @@ def _flat_unflatten(
 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=']')

brainstate/util/_pretty_repr.py

@@ -21,9 +21,10 @@ import dataclasses
 import threading
 from abc import ABC, abstractmethod
 from functools import partial
-from typing import Any, Iterator, Mapping, TypeVar, Union, Callable, Optional
+from typing import Any, Iterator, Mapping, TypeVar, Union, Callable, Optional, Sequence
 
 __all__ = [
+    'yield_unique_pretty_repr_items',
     'PrettyType',
     'PrettyAttr',
     'PrettyRepr',
@@ -80,10 +81,37 @@ class PrettyRepr(ABC):
 
     def __repr__(self) -> str:
         # repr the individual object with the pretty representation
-        return get_repr(self)
+        return pretty_repr_object(self)
 
 
-def _repr_elem(obj: PrettyType, elem: Any) -> str:
+def pretty_repr_elem(obj: PrettyType, elem: Any) -> str:
+    """
+    Constructs a string representation of a single element within a pretty representation.
+
+    This function takes a `PrettyType` object and an element, which must be an instance
+    of `PrettyAttr`, and generates a formatted string that represents the element. The
+    formatting is based on the configuration provided by the `PrettyType` object.
+
+    Parameters
+    ----------
+    obj : PrettyType
+        The configuration object that defines how the element should be formatted.
+        It includes details such as indentation, separators, and surrounding characters.
+    elem : Any
+        The element to be represented. It must be an instance of `PrettyAttr`, which
+        contains the key and value to be formatted.
+
+    Returns
+    -------
+    str
+        A string that represents the element in a formatted manner, adhering to the
+        configuration specified by the `PrettyType` object.
+
+    Raises
+    ------
+    TypeError
+        If the provided element is not an instance of `PrettyAttr`.
+    """
     if not isinstance(elem, PrettyAttr):
         raise TypeError(f'Item must be Elem, got {type(elem).__name__}')
 
@@ -93,9 +121,32 @@ def _repr_elem(obj: PrettyType, elem: Any) -> str:
     return f'{obj.elem_indent}{elem.start}{elem.key}{obj.value_sep}{value}{elem.end}'
 
 
-def get_repr(obj: PrettyRepr) -> str:
+def pretty_repr_object(obj: PrettyRepr) -> str:
     """
-    Get the pretty representation of an object.
+    Generates a pretty string representation of an object that implements the PrettyRepr interface.
+
+    This function utilizes the __pretty_repr__ method of the PrettyRepr interface to obtain
+    a structured representation of the object, which includes both the type and attributes
+    of the object in a human-readable format.
+
+    Parameters
+    ----------
+    obj : PrettyRepr
+        The object for which the pretty representation is to be generated. The object must
+        implement the PrettyRepr interface.
+
+    Returns
+    -------
+    str
+        A string that represents the object in a pretty format, including its type and attributes.
+        The format is determined by the PrettyType and PrettyAttr instances yielded by the
+        __pretty_repr__ method of the object.
+
+    Raises
+    ------
+    TypeError
+        If the provided object does not implement the PrettyRepr interface or if the first item
+        yielded by the __pretty_repr__ method is not an instance of PrettyType.
     """
     if not isinstance(obj, PrettyRepr):
         raise TypeError(f'Object {obj!r} is not representable')
@@ -108,7 +159,7 @@ def get_repr(obj: PrettyRepr) -> str:
         raise TypeError(f'First item must be PrettyType, got {type(obj_repr).__name__}')
 
     # repr attributes
-    elem_reprs = tuple(map(partial(_repr_elem, obj_repr), iterator))
+    elem_reprs = tuple(map(partial(pretty_repr_elem, obj_repr), iterator))
     elems = ',\n'.join(elem_reprs)
     if elems:
         elems = '\n' + elems + '\n'
@@ -140,9 +191,10 @@ class PrettyMapping(PrettyRepr):
     Pretty representation of a mapping.
     """
     mapping: Mapping
+    type_name: str = ''
 
     def __pretty_repr__(self):
-        yield PrettyType(type='', value_sep=': ', start='{', end='}')
+        yield PrettyType(type=self.type_name, value_sep=': ', start='{', end='}')
 
         for key, value in self.mapping.items():
             yield PrettyAttr(repr(key), value)
@@ -150,7 +202,20 @@ class PrettyMapping(PrettyRepr):
 
 @dataclasses.dataclass
 class PrettyReprContext(threading.local):
-    # seen_modules_repr: set[int] | None = None
+    """
+    A thread-local context for managing the state of pretty representation.
+
+    This class is used to keep track of objects that have been seen during
+    the generation of pretty representations, preventing infinite recursion
+    in cases of circular references.
+
+    Attributes
+    ----------
+    seen_modules_repr : dict[int, Any] | None
+        A dictionary mapping object IDs to objects that have been seen
+        during the pretty representation process. This is used to avoid
+        representing the same object multiple times.
+    """
     seen_modules_repr: dict[int, Any] | None = None
 
 
@@ -158,23 +223,80 @@ CONTEXT = PrettyReprContext()
 
 
 def _default_repr_object(node):
+    """
+    Generates a default pretty representation for an object.
+
+    This function yields a `PrettyType` instance that represents the type
+    of the given object. It is used as a default method for representing
+    objects when no custom representation function is provided.
+
+    Parameters
+    ----------
+    node : Any
+        The object for which the pretty representation is to be generated.
+
+    Yields
+    ------
+    PrettyType
+        An instance of `PrettyType` that contains the type information of
+        the object.
+    """
     yield PrettyType(type=type(node))
 
 
 def _default_repr_attr(node):
+    """
+    Generates a default pretty representation for the attributes of an object.
+
+    This function iterates over the attributes of the given object and yields
+    a `PrettyAttr` instance for each attribute that does not start with an
+    underscore. The `PrettyAttr` instances contain the attribute name and its
+    string representation.
+
+    Parameters
+    ----------
+    node : Any
+        The object whose attributes are to be represented.
+
+    Yields
+    ------
+    PrettyAttr
+        An instance of `PrettyAttr` for each non-private attribute of the object,
+        containing the attribute name and its string representation.
+    """
     for name, value in vars(node).items():
         if name.startswith('_'):
             continue
         yield PrettyAttr(name, repr(value))
 
 
-def pretty_repr_avoid_duplicate(
+def yield_unique_pretty_repr_items(
     node,
     repr_object: Optional[Callable] = None,
     repr_attr: Optional[Callable] = None
 ):
     """
-    Pretty representation of an object avoiding duplicate representations.
+    Generates a pretty representation of an object while avoiding duplicate representations.
+
+    This function is designed to yield a structured representation of an object, 
+    using custom or default methods for representing the object itself and its attributes. 
+    It ensures that each object is only represented once to prevent infinite recursion 
+    in cases of circular references.
+
+    Parameters:
+    node : Any
+        The object to be represented.
+    repr_object : Optional[Callable], optional
+        A callable that yields the representation of the object itself. 
+        If not provided, a default representation function is used.
+    repr_attr : Optional[Callable], optional
+        A callable that yields the representation of the object's attributes. 
+        If not provided, a default attribute representation function is used.
+
+    Yields:
+    Union[PrettyType, PrettyAttr]
+        The pretty representation of the object and its attributes, 
+        avoiding duplicates by tracking seen objects.
     """
     if repr_object is None:
         repr_object = _default_repr_object
@@ -206,3 +328,4 @@ def pretty_repr_avoid_duplicate(
     finally:
         if clear_seen:
             CONTEXT.seen_modules_repr = None
+

{brainstate-0.1.0.post20250210.dist-info → brainstate-0.1.0.post20250212.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.1.0.post20250210
+Version: 0.1.0.post20250212
 Summary: A ``State``-based Transformation System for Program Compilation and Augmentation.
 Home-page: https://github.com/chaobrain/brainstate
 Author: BrainState Developers

{brainstate-0.1.0.post20250210.dist-info → brainstate-0.1.0.post20250212.dist-info}/RECORD

@@ -1,12 +1,12 @@
 brainstate/__init__.py,sha256=AkZyyFkn4fB8g2aT6Rc2MO1xICPpUZuDtdze-eUQNc0,1496
-brainstate/_state.py,sha256=Ol-FqHWQnIKmylXHjdsY5izKQhIb0bUw3_UL-7zj4WA,29447
+brainstate/_state.py,sha256=aM2UTfFGvfXfM-pCLvufhgyuuLBGfogBYsz7ZCU8P7Q,28588
 brainstate/_state_test.py,sha256=rJUFRSXEqrrl4qANRewY9mnDlzSbtHwBIGeZ0ku-8Dg,1650
 brainstate/_utils.py,sha256=uJ6WWKq3yb05ZdktCQGLWOXsOJveL1H9pR7eev70Jes,1693
 brainstate/environ.py,sha256=PZnVFWPioUBuWmwCO8wwCKrHQfP3BR-5lYPRl5i5GDA,17698
 brainstate/environ_test.py,sha256=QD6sPCKNtqemVCGwkdImjMazatrvvLr6YeAVcfUnVVY,2045
 brainstate/mixin.py,sha256=g7uVUwZphZWsNs9pb48ozG2cDGaj0hs0g3lq8tDk-Sg,11310
 brainstate/mixin_test.py,sha256=Oq_0fwC9vpXDN4t4dTBhWzLdFDNlcYsrcip14F1yECI,3079
-brainstate/surrogate.py,sha256=t4SzVwUVMAPtC-O1vFbuE9F4265wgAv7ud77ufIJsuk,48464
+brainstate/surrogate.py,sha256=xS4UG4LHKUJdHqwZ5-p-9Y2jWXMa-ssdZJCMiW9zi5k,53540
 brainstate/transform.py,sha256=cxbymTlJ6uHvJWEEYXzFUkAySs_TbUTHakt0NQgWJ3s,808
 brainstate/typing.py,sha256=Qh-LBzm6oG4rSXv4V5qB8SNYcoOR7bASoK_iQxnlafk,10467
 brainstate/augment/__init__.py,sha256=zGPq1eTB_56GRCNC9TiPLKTw07PA2O0OCi7bgjYIrY4,1193
@@ -30,7 +30,7 @@ brainstate/compile/_loop_collect_return.py,sha256=TrKBZhtQecTtuiVz_HOeyepde-znzj
 brainstate/compile/_loop_collect_return_test.py,sha256=bA-_11E8A_0jR5umEO3e409y7bb5QYDTgSL-SBaX7kQ,1802
 brainstate/compile/_loop_no_collection.py,sha256=qto2__Zt2PJntkjB9AXEgraGLvNUJS483BhCXjJyqv0,7495
 brainstate/compile/_loop_no_collection_test.py,sha256=oStB1CSG_iLp9sHdXd1hJNFvlxbzjck9Iy4sABoJDj4,1419
-brainstate/compile/_make_jaxpr.py,sha256=Q-nwm-ibBN0ube4ZjATp924pUkrXuaeT0XgSstqkI40,33304
+brainstate/compile/_make_jaxpr.py,sha256=Rr36U0s8qow1A4KJYXkALX10Rm2pkSYF2j_1eiSuSGI,33292
 brainstate/compile/_make_jaxpr_test.py,sha256=3gwdiutn_PJyiweu3oPEXumxEVHKaE2xDGvkwZy2GEo,4367
 brainstate/compile/_progress_bar.py,sha256=5pCMCEmbTO5XmKtzRUJGA178tuBznWKuh9Kw00wAL1I,7524
 brainstate/compile/_unvmap.py,sha256=CJA6D9lUcBfvdLrpFVvC2AdTJqe9uY0Ht6PltQJyr4U,4228
@@ -42,7 +42,7 @@ brainstate/functional/_normalization.py,sha256=i2EV7hSsqcNdcYRX2wAxjq8doHwyN9eNJ
 brainstate/functional/_others.py,sha256=_u_Ys-LiLzDAP4zJggVwaVvirgoS3jvhXMREoS6JOkM,1737
 brainstate/functional/_spikes.py,sha256=QY-2ayJkgkGELcq-bftPEaf_hJptVf_SP3fY36QvlZc,2678
 brainstate/graph/__init__.py,sha256=noo4TjBg6iEhjjwk0sAGUhR7Ge-z8Vnc2rLYUvnqttw,1295
-brainstate/graph/_graph_node.py,sha256=swAokZLKswSTaq2WEhyLIs38sy_67C6maHI6T3e1hvY,8339
+brainstate/graph/_graph_node.py,sha256=XwzOuaZG9x4eZknQjzJoTnnYAy7wcKD5Vox1VkYr8GM,8345
 brainstate/graph/_graph_node_test.py,sha256=BFGfdzZFDHI0XK7hHotSVWKt3em1taGvn8FHF9NCXx8,2702
 brainstate/graph/_graph_operation.py,sha256=UtBNP7hvxa-5i99LQJStXbFhUbl3icdfTq1oF4MeH1g,64106
 brainstate/graph/_graph_operation_test.py,sha256=zjvpKjQAFWtw8YZuqOk_jmlZNb_-E8oPyNx57dyc8jI,18556
@@ -57,7 +57,7 @@ brainstate/nn/__init__.py,sha256=rxURT8J1XfBn3Vh3Dx_WzVADWn9zVriIty5KZEG-x6o,162
 brainstate/nn/_collective_ops.py,sha256=sSjIIs1MvZA30XFFmK7iL1D_sCeh7hFd3PanCH6kgZo,6779
 brainstate/nn/_exp_euler.py,sha256=yjkfSllFxGWKEAlHo5AzBizzkFj6FEVDKmFV6E2g214,3521
 brainstate/nn/_exp_euler_test.py,sha256=clwRD8QR71k1jn6NrACMDEUcFMh0J9RTosoPnlYWUkw,1242
-brainstate/nn/_module.py,sha256=HDLPvLfB7jat2VT3gBu0MxA7vfzK7xgowemitHX8Cgo,10835
+brainstate/nn/_module.py,sha256=UzsnaTDh5F6Z8B7ou4RXmTdAWXbNkjf03bYP0kF3_fE,10872
 brainstate/nn/_module_test.py,sha256=V4ZhiY_zYPvArkB2eeOTtZcgQrtlRyXKMbS1AJH4vC8,8893
 brainstate/nn/metrics.py,sha256=iupHjSRTHYY-HmEPBC4tXWrZfF4zh1ek2NwSAA0gnwE,14738
 brainstate/nn/_dyn_impl/__init__.py,sha256=Oazar7h89dp1WA2Vx4Tj7gCBhxJKH4LAUEABkBEG7vU,1462
@@ -84,7 +84,7 @@ brainstate/nn/_elementwise/_dropout_test.py,sha256=k6aB5v8RYMoV5w8UV9UNSFhaQTV7w
 brainstate/nn/_elementwise/_elementwise.py,sha256=om-KpwDTk5yFG5KBYXXHquRLV7s28_FJjk-omvyMyvQ,33342
 brainstate/nn/_elementwise/_elementwise_test.py,sha256=SZI9jB39sZ5SO1dpWGW-PhodthwN0GU9FY1nqf2fWcs,5341
 brainstate/nn/_interaction/__init__.py,sha256=TTY_SeNrdx4VnUSw6vdyl02OHdS9Qs15cWBp6kjsyNQ,1289
-brainstate/nn/_interaction/_conv.py,sha256=lwyxTVsJVPiKlZcgB6iqE64aX7AOJzplDSj4y6-m18o,18592
+brainstate/nn/_interaction/_conv.py,sha256=eKhABWtG3QlOy7TPY9yoQjP3liBh9bb8X5Wns3_YUUQ,18499
 brainstate/nn/_interaction/_conv_test.py,sha256=fHXRFYnDghFiKre63RqMwIE_gbPKdK34UPhKOz-J3qU,8695
 brainstate/nn/_interaction/_embedding.py,sha256=iK0I1ExKWFa_QzV9UDGj32Ljsmdr1g_LlAtMcusebxU,2187
 brainstate/nn/_interaction/_linear.py,sha256=EnkOk1oE79rvRIjU6HBllxUpVOEcQQCj4vtavo9AJjI,14767
@@ -109,16 +109,16 @@ brainstate/random/_rand_state.py,sha256=nuoQ8GU1MfJPRNN-ZmRQsggVjoyPhaEdZmwM7_4-
 brainstate/random/_random_for_unit.py,sha256=kGp4EUX19MXJ9Govoivbg8N0bddqOldKEI2h_TbdONY,2057
 brainstate/util/__init__.py,sha256=-FWEuSKXG3mWxYphGFAy3UEuVe39lFs1GruluzdXDoI,1502
 brainstate/util/_caller.py,sha256=T3bzu7-09r-6EOrU6Muca_aMXSQua_X2lXjEqb-w39w,2782
-brainstate/util/_dict.py,sha256=Yapug-_RZQYjvd8cZ3v90_MX7rUYJDBzBnZJT6a0NXY,26178
+brainstate/util/_dict.py,sha256=qPUbqjRVHUvVHhSWBPojx_srsh6-iy1k5oPMn1DdrnQ,30880
 brainstate/util/_dict_test.py,sha256=Dn0TdjX6wLBXaTD4jfYTu6cKfFHwKSxi4_3bX7kB_IA,5621
 brainstate/util/_error.py,sha256=eyZ8PGFixqe2K5OEfjSDzI-2tU0ieYQoUpBP7yStlPQ,878
 brainstate/util/_filter.py,sha256=1-bvFHdjeehvXeHTrCEp8xr25lopKe8d3XZGCNegq0s,4970
 brainstate/util/_others.py,sha256=jsPZwP-v_5HRV-LB5F0NUsiqr04y8bmGIsu_JMyVcbQ,14762
-brainstate/util/_pretty_repr.py,sha256=bDpU4gbkS4B8cXBkiN8kBQNmruxiJzDRF-eIqzyeYnM,5716
+brainstate/util/_pretty_repr.py,sha256=-TZPIgfTLB-Eg7rgT7KAkV1r-HX0q6nCgKDKA7Qdsw4,10577
 brainstate/util/_scaling.py,sha256=pc_eM_SZVwkY65I4tJh1ODiHNCoEhsfFXl2zBK0PLAg,7562
 brainstate/util/_struct.py,sha256=KMMHcshOM20gYhSahNzWLxsTt-Rt3AeX3Uz26-rP9vI,17619
-brainstate-0.1.0.post20250210.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.1.0.post20250210.dist-info/METADATA,sha256=__N9QGz8FFW6rXXG3_Y5YTKFd9iWM_MVddxJP74hZcI,3585
-brainstate-0.1.0.post20250210.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.1.0.post20250210.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.1.0.post20250210.dist-info/RECORD,,
+brainstate-0.1.0.post20250212.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.1.0.post20250212.dist-info/METADATA,sha256=OVPO4wr0e0j_Lvk_OQKTpTdNUbOGsFt_BW_qKakO8xE,3585
+brainstate-0.1.0.post20250212.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+brainstate-0.1.0.post20250212.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.1.0.post20250212.dist-info/RECORD,,