brainstate 0.1.0.post20250206__py2.py3-none-any.whl → 0.1.0.post20250209__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, PrettyReprTree
 
 __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], PrettyReprTree):
     """
     The pointer to specify the dynamical data.
 
@@ -261,8 +263,9 @@ class State(Generic[A], PrettyRepr):
         """
         The data and its value.
         """
+        self.check_if_deleted()
         record_state_value_read(self)
-        return self._read_value()
+        return self._value
 
     @value.setter
     def value(self, v) -> None:
@@ -272,14 +275,7 @@ class State(Generic[A], PrettyRepr):
         Args:
           v: The value.
         """
-        # NOTE: the following order is important
-
-        if isinstance(v, State):  # value checking
-            raise ValueError('Cannot set value to a State, ' 'use `copy_from` method instead')
-        self._check_value_tree(v)  # check the tree structure
-        record_state_value_write(self)  # record the value by the stack (>= level)
-        self._been_writen = True  # set the flag
-        self._write_value(v)  # write the value
+        self.write_value(v)
 
     @property
     def stack_level(self):
@@ -301,18 +297,17 @@ class State(Generic[A], PrettyRepr):
         """
         self._level = level
 
-    def _read_value(self) -> PyTree[ArrayLike]:
-        """
-        The interface to customize the value reading.
-        """
-        self.check_if_deleted()
-        return self._value
-
-    def _write_value(self, v) -> None:
-        """
-        The interface to customize the value writing.
-        """
+    def write_value(self, v) -> None:
+        # value checking
+        if isinstance(v, State):
+            raise ValueError('Cannot set value to a State, ' 'use `copy_from` method instead')
+        self._check_value_tree(v)
+        # write the value by the stack (>= level)
+        record_state_value_write(self)
+        # set the value
         self._value = v
+        # set flag
+        self._been_writen = True
 
     def restore_value(self, v) -> None:
         """
@@ -434,45 +429,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)
@@ -809,7 +784,7 @@ class StateTraceStack(Generic[A]):
         return StateTraceStack().merge(self, other)
 
 
-class TreefyState(Generic[A], PrettyRepr):
+class TreefyState(Generic[A], PrettyReprTree):
     """
     The state as a pytree.
     """
@@ -831,35 +806,19 @@ 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':
+            if self.name is None:
+                return None, None
+            else:
+                return 'name', v
+
+        return k, v
 
     def replace(self, value: B) -> TreefyState[B]:
         """

brainstate/compile/_jit.py

@@ -82,6 +82,7 @@ def _get_jitted_fun(
             return fun.fun(*args, **params)
 
         # compile the function and get the state trace
+        # print('Compiling ...')
         state_trace = fun.compile_function_and_get_state_trace(*args, **params, return_only_write=True)
         read_state_vals = state_trace.get_read_state_values(True)
 

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}, "
@@ -499,7 +499,7 @@ class StatefulFunction(object):
         state_vals, out = self.jaxpr_call([st.value for st in state_trace.states], *args, **kwargs)
         for st, written, val in zip(state_trace.states, state_trace.been_writen, state_vals):
             if written:
-                st.value = val
+                st.write_value(val)
             else:
                 st.restore_value(val)
         return out

brainstate/compile/_util.py

@@ -31,7 +31,7 @@ def write_back_state_values(
     assert len(state_trace.states) == len(state_trace.been_writen) == len(read_state_vals) == len(write_state_vals)
     for st, write, val_r, val_w in zip(state_trace.states, state_trace.been_writen, read_state_vals, write_state_vals):
         if write:
-            st.value = val_w
+            st.write_value(val_w)
         else:
             st.restore_value(val_r)
 

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/graph/_graph_operation.py

@@ -609,7 +609,7 @@ def _get_children(graph_def, state_mapping, index_ref, index_ref_cache):
                             variable.update_from_ref(value)
                         elif isinstance(value, State):
                             if value._been_writen:
-                                variable.value = value.value
+                                variable.write_value(value.value)
                             else:
                                 variable.restore_value(value.value)
                         else:

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/util/_dict.py

@@ -24,11 +24,12 @@ 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
 from ._struct import dataclass
 
 __all__ = [
-    'NestedDict', 'FlattedDict', 'flat_mapping', 'nest_mapping',
+    'PrettyDict', 'NestedDict', 'FlattedDict', 'flat_mapping', 'nest_mapping',
+    'PrettyList', 'PrettyReprTree',
 ]
 
 A = TypeVar('A')
@@ -40,6 +41,44 @@ ExtractValueFn = abc.Callable[[Any], Any]
 SetValueFn = abc.Callable[[V, Any], V]
 
 
+
+
+class PrettyReprTree(PrettyRepr):
+    """
+    Pretty representation of a tree.
+    """
+
+    def __pretty_repr__(self):
+        return yield_unique_pretty_repr_items(
+            self,
+            repr_object=self._repr_object,
+            repr_attr=self._repr_attr,
+        )
+
+    def __pretty_repr_item__(self, k, v):
+        return k, v
+
+    def _repr_object(self, node: PrettyDict):
+        yield PrettyType(type(node), value_sep=': ', start='({', end='})')
+
+    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
+
+            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)
+
+
 # the empty node is a struct.dataclass to be compatible with JAX.
 @dataclass
 class _EmptyNode:
@@ -213,10 +252,10 @@ class PrettyDict(dict, PrettyRepr):
 
     def __repr__(self) -> str:
         # repr the individual object with the pretty representation
-        return get_repr(self)
+        return pretty_repr(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 +276,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 +779,37 @@ 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(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,11 @@ 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',
+    'pretty_repr',
     'PrettyType',
     'PrettyAttr',
     'PrettyRepr',
@@ -80,7 +82,7 @@ class PrettyRepr(ABC):
 
     def __repr__(self) -> str:
         # repr the individual object with the pretty representation
-        return get_repr(self)
+        return pretty_repr(self)
 
 
 def _repr_elem(obj: PrettyType, elem: Any) -> str:
@@ -93,7 +95,7 @@ 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(obj: PrettyRepr) -> str:
     """
     Get the pretty representation of an object.
     """
@@ -140,9 +142,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)
@@ -168,7 +171,7 @@ def _default_repr_attr(node):
         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
@@ -206,3 +209,4 @@ def pretty_repr_avoid_duplicate(
     finally:
         if clear_seen:
             CONTEXT.seen_modules_repr = None
+

{brainstate-0.1.0.post20250206.dist-info → brainstate-0.1.0.post20250209.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.1.0.post20250206
+Version: 0.1.0.post20250209
 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.post20250206.dist-info → brainstate-0.1.0.post20250209.dist-info}/RECORD

@@ -1,5 +1,5 @@
 brainstate/__init__.py,sha256=AkZyyFkn4fB8g2aT6Rc2MO1xICPpUZuDtdze-eUQNc0,1496
-brainstate/_state.py,sha256=Ol-FqHWQnIKmylXHjdsY5izKQhIb0bUw3_UL-7zj4WA,29447
+brainstate/_state.py,sha256=W1Q_RAL01rUSLZuOARMuX9I-26tBuIl_VzNWAziz6A8,27518
 brainstate/_state_test.py,sha256=rJUFRSXEqrrl4qANRewY9mnDlzSbtHwBIGeZ0ku-8Dg,1650
 brainstate/_utils.py,sha256=uJ6WWKq3yb05ZdktCQGLWOXsOJveL1H9pR7eev70Jes,1693
 brainstate/environ.py,sha256=PZnVFWPioUBuWmwCO8wwCKrHQfP3BR-5lYPRl5i5GDA,17698
@@ -24,17 +24,17 @@ brainstate/compile/_conditions.py,sha256=gApsHKGQrf1QBjoKXDVL7VsoeJ2zFtSc-hFz9nb
 brainstate/compile/_conditions_test.py,sha256=s9LF6h9LvigvgxUIugTqvgCHBIU8TXS1Ar1OlIxXfrw,8389
 brainstate/compile/_error_if.py,sha256=TFvhqITKkRO9m30GdlUP4eEjJvLWQUhjkujXO9zvrWs,2689
 brainstate/compile/_error_if_test.py,sha256=OdJG483IIdOrCHxtHd49OHfOxCSnSkk7GdAUOzSt8bE,2044
-brainstate/compile/_jit.py,sha256=-Y8fyy8gc7qQT2ti4-N-74hjP_6C-D8YC5h-1unEKuI,13910
+brainstate/compile/_jit.py,sha256=itAWENKfJvnlaWl_uSy8lHTK8K1in89F_ZXXwp-EGRM,13944
 brainstate/compile/_jit_test.py,sha256=zD7kck9SQJGmUDolh9P4luKwQ21fBGje1Z4STTEXIuA,4135
 brainstate/compile/_loop_collect_return.py,sha256=TrKBZhtQecTtuiVz_HOeyepde-znzjlyk0Te53-AvOE,23492
 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=MuAa9LjXi29DjYgDUrK0WaomkjbhHZk9mWW04XGcV98,33297
 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
-brainstate/compile/_util.py,sha256=iKk51BHAerBFj2BTxPNdjsk3MZQiXenzpCr7Ys0iYWg,6299
+brainstate/compile/_util.py,sha256=aCvkTV--g4NsqcodTdBAISt4EwgezCbKzNUV58n-Q_Y,6304
 brainstate/functional/__init__.py,sha256=j6-3Er4fgqWpvntzYCZVB3e5hoz-Z3aqvapITCuDri0,1107
 brainstate/functional/_activations.py,sha256=S0Ok7sq5FTbmJWSejpOCHo1jpKX0gYOLy_TO2IUXM8s,21726
 brainstate/functional/_activations_test.py,sha256=T___RlSrIfXwlkw8dg5A9EZMTZGDzv3a2evUwq_nYFg,13034
@@ -42,9 +42,9 @@ 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.py,sha256=cIwGo3ICgtce2fmdn917r81evMFjJIKeW9doaQK4DD8,64111
 brainstate/graph/_graph_operation_test.py,sha256=zjvpKjQAFWtw8YZuqOk_jmlZNb_-E8oPyNx57dyc8jI,18556
 brainstate/init/__init__.py,sha256=R1dHgub47o-WJM9QkFLc7x_Q7GsyaKKDtrRHTFPpC5g,1097
 brainstate/init/_base.py,sha256=B_NLS9aKNrvuj5NAlSgBbQTVev7IRvzcx8vH0J-Gq2w,1671
@@ -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
@@ -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=tb5nPrTKJe4G_BDv33XYTUaYQDz6od-5psG4TKemc7A,28111
 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=vNwRlj4sI4QJ_koyIs7eKdUMeB_QWwzRYsE8PpAWN3g,5833
 brainstate/util/_scaling.py,sha256=pc_eM_SZVwkY65I4tJh1ODiHNCoEhsfFXl2zBK0PLAg,7562
 brainstate/util/_struct.py,sha256=KMMHcshOM20gYhSahNzWLxsTt-Rt3AeX3Uz26-rP9vI,17619
-brainstate-0.1.0.post20250206.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.1.0.post20250206.dist-info/METADATA,sha256=-H1vqEufSWo4cx5wCLeYECZuStnZXDZ_Vd9JBL5dni4,3585
-brainstate-0.1.0.post20250206.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.1.0.post20250206.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.1.0.post20250206.dist-info/RECORD,,
+brainstate-0.1.0.post20250209.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.1.0.post20250209.dist-info/METADATA,sha256=vc9kKmrq5JM9Os6brL4zecy55nEpd9ASK9GZNJBQV9g,3585
+brainstate-0.1.0.post20250209.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+brainstate-0.1.0.post20250209.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.1.0.post20250209.dist-info/RECORD,,