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

brainstate/_state.py

@@ -30,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, PrettyReprTree
+from brainstate.util import DictManager, PrettyObject
 
 __all__ = [
     'State', 'ShortTermState', 'LongTermState', 'HiddenState', 'ParamState', 'TreefyState',
@@ -186,7 +186,7 @@ def _get_trace_stack_level() -> int:
     return len(TRACE_CONTEXT.state_stack)
 
 
-class State(Generic[A], PrettyReprTree):
+class State(Generic[A], PrettyObject):
     """
     The pointer to specify the dynamical data.
 
@@ -465,6 +465,25 @@ class State(Generic[A], PrettyReprTree):
         """
         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
@@ -791,7 +810,7 @@ class StateTraceStack(Generic[A]):
         return StateTraceStack().merge(self, other)
 
 
-class TreefyState(Generic[A], PrettyReprTree):
+class TreefyState(Generic[A], PrettyObject):
     """
     The state as a pytree.
     """

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,12 +24,17 @@ import jax
 
 from brainstate.typing import Filter, PathParts
 from ._filter import to_predicate
-from ._pretty_repr import PrettyRepr, PrettyType, PrettyAttr, yield_unique_pretty_repr_items, pretty_repr
+from ._pretty_repr import PrettyRepr, PrettyType, PrettyAttr, yield_unique_pretty_repr_items, pretty_repr_object
 from ._struct import dataclass
 
 __all__ = [
-    'PrettyDict', 'NestedDict', 'FlattedDict', 'flat_mapping', 'nest_mapping',
-    'PrettyList', 'PrettyReprTree',
+    'PrettyDict',
+    'NestedDict',
+    'FlattedDict',
+    'flat_mapping',
+    'nest_mapping',
+    'PrettyList',
+    'PrettyObject',
 ]
 
 A = TypeVar('A')
@@ -41,42 +46,117 @@ 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.
 
-class PrettyReprTree(PrettyRepr):
+    Yields:
+        PrettyType: A PrettyType object representing the type of the node,
+        with specified value separator, start, and end characters.
     """
-    Pretty representation of a tree.
+    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):
-        return yield_unique_pretty_repr_items(
+        """
+        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=self._repr_object,
-            repr_attr=self._repr_attr,
+            repr_object=_repr_object_general,
+            repr_attr=_repr_attribute_general,
         )
 
     def __pretty_repr_item__(self, k, v):
-        return k, v
-
-    def _repr_object(self, node: PrettyDict):
-        yield PrettyType(type(node), value_sep=': ', start='({', end='})')
+        """
+        Returns a tuple of the key and value for pretty representation.
 
-    def _repr_attr(self, node):
-        for k, v in vars(node).items():
-            k, v = self.__pretty_repr_item__(k, v)
-            if k is None:
-                continue
+        This method is used to generate a pretty representation of an item
+        in a data structure, typically for debugging or logging purposes.
 
-            if isinstance(v, list):
-                v = PrettyList(v)
+        Args:
+            k: The key of the item.
+            v: The value of the item.
 
-            if isinstance(v, dict):
-                v = PrettyDict(v)
+        Returns:
+            A tuple containing the key and value.
+        """
+        return k, v
 
-            if isinstance(v, PrettyDict):
-                v = NestedStateRepr(v)
 
-            yield PrettyAttr(repr(k), v)
+PrettyReprTree = PrettyObject
 
 
 # the empty node is a struct.dataclass to be compatible with JAX.
@@ -252,7 +332,7 @@ class PrettyDict(dict, PrettyRepr):
 
     def __repr__(self) -> str:
         # repr the individual object with the pretty representation
-        return pretty_repr(self)
+        return pretty_repr_object(self)
 
     def __pretty_repr__(self):
         yield from yield_unique_pretty_repr_items(self, _default_repr_object, _default_repr_attr)
@@ -789,7 +869,7 @@ class PrettyList(list, PrettyRepr):
         yield from yield_unique_pretty_repr_items(self, _list_repr_object, _list_repr_attr)
 
     def __repr__(self):
-        return pretty_repr(self)
+        return pretty_repr_object(self)
 
     def tree_flatten(self):
         return list(self), ()
@@ -812,4 +892,3 @@ def _list_repr_attr(node: PrettyList):
 
 def _list_repr_object(node: PrettyDict):
     yield PrettyType('', value_sep='', start='[', end=']')
-

brainstate/util/_pretty_repr.py

@@ -25,7 +25,6 @@ from typing import Any, Iterator, Mapping, TypeVar, Union, Callable, Optional, S
 
 __all__ = [
     'yield_unique_pretty_repr_items',
-    'pretty_repr',
     'PrettyType',
     'PrettyAttr',
     'PrettyRepr',
@@ -82,10 +81,37 @@ class PrettyRepr(ABC):
 
     def __repr__(self) -> str:
         # repr the individual object with the pretty representation
-        return pretty_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__}')
 
@@ -95,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 pretty_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')
@@ -110,7 +159,7 @@ def pretty_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'
@@ -153,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
 
 
@@ -161,10 +223,47 @@ 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
@@ -177,7 +276,27 @@ def yield_unique_pretty_repr_items(
     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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.1.0.post20250211
+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.post20250211.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=TrlgaX_Hu3aSdkOUqw88HZMKUKM_PgnE9YGp595nfi0,27761
+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
@@ -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=tb5nPrTKJe4G_BDv33XYTUaYQDz6od-5psG4TKemc7A,28111
+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=vNwRlj4sI4QJ_koyIs7eKdUMeB_QWwzRYsE8PpAWN3g,5833
+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.post20250211.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.1.0.post20250211.dist-info/METADATA,sha256=werv_oEsECW5xLgvO__Yjth5vLfz-YARh442Q6E6FIk,3585
-brainstate-0.1.0.post20250211.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.1.0.post20250211.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.1.0.post20250211.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,,