brainstate 0.1.4__py2.py3-none-any.whl → 0.1.6__py2.py3-none-any.whl

brainstate/util/pretty_pytree.py

@@ -373,19 +373,19 @@ class NestedDict(PrettyDict):
 
         Example usage::
 
-          >>> import brainstate as bst
+          >>> import brainstate as brainstate
 
-          >>> class Model(bst.nn.Module):
+          >>> class Model(brainstate.nn.Module):
           ...   def __init__(self):
           ...     super().__init__()
-          ...     self.batchnorm = bst.nn.BatchNorm1d([10, 3])
-          ...     self.linear = bst.nn.Linear(2, 3)
+          ...     self.batchnorm = brainstate.nn.BatchNorm1d([10, 3])
+          ...     self.linear = brainstate.nn.Linear(2, 3)
           ...   def __call__(self, x):
           ...     return self.linear(self.batchnorm(x))
 
           >>> model = Model()
-          >>> state_map = bst.graph.treefy_states(model)
-          >>> param, others = state_map.treefy_split(bst.ParamState, ...)
+          >>> state_map = brainstate.graph.treefy_states(model)
+          >>> param, others = state_map.treefy_split(brainstate.ParamState, ...)
 
         Arguments:
           first: The first filter
@@ -495,14 +495,14 @@ class FlattedDict(PrettyDict):
 
     Example usage::
 
-      >>> import brainstate as bst
+      >>> import brainstate as brainstate
       >>> import jax.numpy as jnp
       >>>
-      >>> class Model(bst.nn.Module):
+      >>> class Model(brainstate.nn.Module):
       ...   def __init__(self):
       ...     super().__init__()
-      ...     self.batchnorm = bst.nn.BatchNorm1d([10, 3])
-      ...     self.linear = bst.nn.Linear(2, 3)
+      ...     self.batchnorm = brainstate.nn.BatchNorm1d([10, 3])
+      ...     self.linear = brainstate.nn.Linear(2, 3)
       ...   def __call__(self, x):
       ...     return self.linear(self.batchnorm(x))
       >>>

brainstate/util/{_pretty_pytree_test.py → pretty_pytree_test.py}

@@ -13,31 +13,30 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
 
 import unittest
 
 import jax
 from absl.testing import absltest
 
-import brainstate as bst
+import brainstate
 
 
 class TestNestedMapping(absltest.TestCase):
     def test_create_state(self):
-        state = bst.util.NestedDict({'a': bst.ParamState(1), 'b': {'c': bst.ParamState(2)}})
+        state = brainstate.util.NestedDict({'a': brainstate.ParamState(1), 'b': {'c': brainstate.ParamState(2)}})
 
         assert state['a'].value == 1
         assert state['b']['c'].value == 2
 
     def test_get_attr(self):
-        state = bst.util.NestedDict({'a': bst.ParamState(1), 'b': {'c': bst.ParamState(2)}})
+        state = brainstate.util.NestedDict({'a': brainstate.ParamState(1), 'b': {'c': brainstate.ParamState(2)}})
 
         assert state.a.value == 1
         assert state.b['c'].value == 2
 
     def test_set_attr(self):
-        state = bst.util.NestedDict({'a': bst.ParamState(1), 'b': {'c': bst.ParamState(2)}})
+        state = brainstate.util.NestedDict({'a': brainstate.ParamState(1), 'b': {'c': brainstate.ParamState(2)}})
 
         state.a.value = 3
         state.b['c'].value = 4
@@ -46,36 +45,36 @@ class TestNestedMapping(absltest.TestCase):
         assert state['b']['c'].value == 4
 
     def test_set_attr_variables(self):
-        state = bst.util.NestedDict({'a': bst.ParamState(1), 'b': {'c': bst.ParamState(2)}})
+        state = brainstate.util.NestedDict({'a': brainstate.ParamState(1), 'b': {'c': brainstate.ParamState(2)}})
 
         state.a.value = 3
         state.b['c'].value = 4
 
-        assert isinstance(state.a, bst.ParamState)
+        assert isinstance(state.a, brainstate.ParamState)
         assert state.a.value == 3
-        assert isinstance(state.b['c'], bst.ParamState)
+        assert isinstance(state.b['c'], brainstate.ParamState)
         assert state.b['c'].value == 4
 
     def test_add_nested_attr(self):
-        state = bst.util.NestedDict({'a': bst.ParamState(1), 'b': {'c': bst.ParamState(2)}})
-        state.b['d'] = bst.ParamState(5)
+        state = brainstate.util.NestedDict({'a': brainstate.ParamState(1), 'b': {'c': brainstate.ParamState(2)}})
+        state.b['d'] = brainstate.ParamState(5)
 
         assert state['b']['d'].value == 5
 
     def test_delete_nested_attr(self):
-        state = bst.util.NestedDict({'a': bst.ParamState(1), 'b': {'c': bst.ParamState(2)}})
+        state = brainstate.util.NestedDict({'a': brainstate.ParamState(1), 'b': {'c': brainstate.ParamState(2)}})
         del state['b']['c']
 
         assert 'c' not in state['b']
 
     def test_integer_access(self):
-        class Foo(bst.nn.Module):
+        class Foo(brainstate.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.layers = [bst.nn.Linear(1, 2), bst.nn.Linear(2, 3)]
+                self.layers = [brainstate.nn.Linear(1, 2), brainstate.nn.Linear(2, 3)]
 
         module = Foo()
-        state_refs = bst.graph.treefy_states(module)
+        state_refs = brainstate.graph.treefy_states(module)
 
         assert module.layers[0].weight.value['weight'].shape == (1, 2)
         assert state_refs.layers[0]['weight'].value['weight'].shape == (1, 2)
@@ -83,8 +82,8 @@ class TestNestedMapping(absltest.TestCase):
         assert state_refs.layers[1]['weight'].value['weight'].shape == (2, 3)
 
     def test_pure_dict(self):
-        module = bst.nn.Linear(4, 5)
-        state_map = bst.graph.treefy_states(module)
+        module = brainstate.nn.Linear(4, 5)
+        state_map = brainstate.graph.treefy_states(module)
         pure_dict = state_map.to_pure_dict()
         assert isinstance(pure_dict, dict)
         assert isinstance(pure_dict['weight'].value['weight'], jax.Array)
@@ -93,27 +92,27 @@ class TestNestedMapping(absltest.TestCase):
 
 class TestSplit(unittest.TestCase):
     def test_split(self):
-        class Model(bst.nn.Module):
+        class Model(brainstate.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.batchnorm = bst.nn.BatchNorm1d([10, 3])
-                self.linear = bst.nn.Linear([10, 3], [10, 4])
+                self.batchnorm = brainstate.nn.BatchNorm1d([10, 3])
+                self.linear = brainstate.nn.Linear([10, 3], [10, 4])
 
             def __call__(self, x):
                 return self.linear(self.batchnorm(x))
 
-        with bst.environ.context(fit=True):
+        with brainstate.environ.context(fit=True):
             model = Model()
-            x = bst.random.randn(1, 10, 3)
+            x = brainstate.random.randn(1, 10, 3)
             y = model(x)
             self.assertEqual(y.shape, (1, 10, 4))
 
-        state_map = bst.graph.treefy_states(model)
+        state_map = brainstate.graph.treefy_states(model)
 
         with self.assertRaises(ValueError):
-            params, others = state_map.split(bst.ParamState)
+            params, others = state_map.split(brainstate.ParamState)
 
-        params, others = state_map.split(bst.ParamState, ...)
+        params, others = state_map.split(brainstate.ParamState, ...)
         print()
         print(params)
         print(others)
@@ -124,37 +123,37 @@ class TestSplit(unittest.TestCase):
 
 class TestStateMap2(unittest.TestCase):
     def test1(self):
-        class Model(bst.nn.Module):
+        class Model(brainstate.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.batchnorm = bst.nn.BatchNorm1d([10, 3])
-                self.linear = bst.nn.Linear([10, 3], [10, 4])
+                self.batchnorm = brainstate.nn.BatchNorm1d([10, 3])
+                self.linear = brainstate.nn.Linear([10, 3], [10, 4])
 
             def __call__(self, x):
                 return self.linear(self.batchnorm(x))
 
-        with bst.environ.context(fit=True):
+        with brainstate.environ.context(fit=True):
             model = Model()
-            state_map = bst.graph.treefy_states(model).to_flat()
-            state_map = bst.util.NestedDict(state_map)
+            state_map = brainstate.graph.treefy_states(model).to_flat()
+            state_map = brainstate.util.NestedDict(state_map)
 
 
 class TestFlattedMapping(unittest.TestCase):
     def test1(self):
-        class Model(bst.nn.Module):
+        class Model(brainstate.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.batchnorm = bst.nn.BatchNorm1d([10, 3])
-                self.linear = bst.nn.Linear([10, 3], [10, 4])
+                self.batchnorm = brainstate.nn.BatchNorm1d([10, 3])
+                self.linear = brainstate.nn.Linear([10, 3], [10, 4])
 
             def __call__(self, x):
                 return self.linear(self.batchnorm(x))
 
         model = Model()
         # print(model.states())
-        # print(bst.graph.states(model))
-        self.assertTrue(model.states() == bst.graph.states(model))
+        # print(brainstate.graph.states(model))
+        self.assertTrue(model.states() == brainstate.graph.states(model))
 
         print(model.nodes())
-        # print(bst.graph.nodes(model))
-        self.assertTrue(model.nodes() == bst.graph.nodes(model))
+        # print(brainstate.graph.nodes(model))
+        self.assertTrue(model.nodes() == brainstate.graph.nodes(model))

brainstate/util/struct.py

@@ -56,16 +56,16 @@ def dataclass(clz: T, **kwargs) -> T:
     The ``dataclass`` decorator makes it easy to define custom classes that can be
     passed safely to Jax. For example::
 
-      >>> import brainstate as bst
+      >>> import brainstate as brainstate
       >>> import jax
       >>> from typing import Any, Callable
 
-      >>> @bst.util.dataclass
+      >>> @brainstate.util.dataclass
       ... class Model:
       ...   params: Any
       ...   # use pytree_node=False to indicate an attribute should not be touched
       ...   # by Jax transformations.
-      ...   apply_fn: Callable = bst.util.field(pytree_node=False)
+      ...   apply_fn: Callable = brainstate.util.field(pytree_node=False)
 
       ...   def __apply__(self, *args):
       ...     return self.apply_fn(*args)
@@ -97,7 +97,7 @@ def dataclass(clz: T, **kwargs) -> T:
     This way the simple constructor used by ``jax.tree_util`` is
     preserved. Consider the following example::
 
-      >>> @bst.util.dataclass
+      >>> @brainstate.util.dataclass
       ... class DirectionAndScaleKernel:
       ...   direction: jax.Array
       ...   scale: jax.Array
@@ -189,15 +189,15 @@ class PyTreeNode:
 
     Example::
 
-      >>> import brainstate as bst
+      >>> import brainstate as brainstate
       >>> import jax
       >>> from typing import Any, Callable
 
-      >>> class Model(bst.util.PyTreeNode):
+      >>> class Model(brainstate.util.PyTreeNode):
       ...   params: Any
       ...   # use pytree_node=False to indicate an attribute should not be touched
       ...   # by Jax transformations.
-      ...   apply_fn: Callable = bst.util.field(pytree_node=False)
+      ...   apply_fn: Callable = brainstate.util.field(pytree_node=False)
 
       ...   def __apply__(self, *args):
       ...     return self.apply_fn(*args)

{brainstate-0.1.4.dist-info → brainstate-0.1.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.1.4
+Version: 0.1.6
 Summary: A ``State``-based Transformation System for Program Compilation and Augmentation.
 Home-page: https://github.com/chaobrain/brainstate
 Author: BrainState Developers
@@ -31,22 +31,22 @@ License-File: LICENSE
 Requires-Dist: jax
 Requires-Dist: jaxlib
 Requires-Dist: numpy
-Requires-Dist: brainunit (>=0.0.4)
+Requires-Dist: brainunit>=0.0.4
 Requires-Dist: brainevent
 Provides-Extra: cpu
-Requires-Dist: jax[cpu] ; extra == 'cpu'
-Requires-Dist: brainunit[cpu] ; extra == 'cpu'
-Requires-Dist: brainevent[cpu] ; extra == 'cpu'
+Requires-Dist: jax[cpu]; extra == "cpu"
+Requires-Dist: brainunit[cpu]; extra == "cpu"
+Requires-Dist: brainevent[cpu]; extra == "cpu"
 Provides-Extra: cuda12
-Requires-Dist: jax[cuda12] ; extra == 'cuda12'
-Requires-Dist: brainunit[cuda12] ; extra == 'cuda12'
-Requires-Dist: brainevent[cuda12] ; extra == 'cuda12'
+Requires-Dist: jax[cuda12]; extra == "cuda12"
+Requires-Dist: brainunit[cuda12]; extra == "cuda12"
+Requires-Dist: brainevent[cuda12]; extra == "cuda12"
 Provides-Extra: testing
-Requires-Dist: pytest ; extra == 'testing'
+Requires-Dist: pytest; extra == "testing"
 Provides-Extra: tpu
-Requires-Dist: jax[tpu] ; extra == 'tpu'
-Requires-Dist: brainunit[tpu] ; extra == 'tpu'
-Requires-Dist: brainevent[tpu] ; extra == 'tpu'
+Requires-Dist: jax[tpu]; extra == "tpu"
+Requires-Dist: brainunit[tpu]; extra == "tpu"
+Requires-Dist: brainevent[tpu]; extra == "tpu"
 
 
 # A ``State``-based Transformation System for Program Compilation and Augmentation

{brainstate-0.1.4.dist-info → brainstate-0.1.6.dist-info}/RECORD

@@ -1,27 +1,27 @@
-brainstate/__init__.py,sha256=A_OD4cJiVu3xpthNGJh6fhRjCKsI7_Mxsow3Al4m2-w,1496
+brainstate/__init__.py,sha256=jCLuGmUnNNYYpHUGsGZrw--LWrAtO0m1ABtU_1qvtRM,1496
 brainstate/_compatible_import.py,sha256=LUSZlA0APWozxM8Kf9pZrM2YbwY7X3jzVVHZInaBL7Y,4630
-brainstate/_state.py,sha256=o5Kk4bGwVz6Dfj9dlmZqdh6zUXcz6Tvc6WOjH9ajlIU,60716
+brainstate/_state.py,sha256=PLzoYx13jIgnzyBxnktLTVPCFl6seG__aNIHS9A5Nms,60770
 brainstate/_state_test.py,sha256=b6uvZdVRyC4n6-fYzmHNry1b-gJ6zE_kRSxGinqiHaw,1638
 brainstate/_utils.py,sha256=j-b239RHfC5BnvhGbSExQpdY21LrMEyWMSHBdNGThOI,1657
-brainstate/environ.py,sha256=VgtG0S_aR1g_1gplRWg_v2ZrcS-F6LZk35BPCBgsIIA,17660
+brainstate/environ.py,sha256=vxwGyYm4lf9Ym5iD6YHsYy4bTwmeD8yHIuZsqqGfP4s,17793
 brainstate/environ_test.py,sha256=QD6sPCKNtqemVCGwkdImjMazatrvvLr6YeAVcfUnVVY,2045
-brainstate/mixin.py,sha256=LOS8WPfB6OBevdZ9wh_U6qzuw4mrCNAMcJy0yIitIBg,11522
+brainstate/mixin.py,sha256=eOOODmESCYITJbvz6dnQLyu8u2m0YrTaEvZsepdAUto,10991
 brainstate/mixin_test.py,sha256=Oq_0fwC9vpXDN4t4dTBhWzLdFDNlcYsrcip14F1yECI,3079
-brainstate/surrogate.py,sha256=yPc-uhJl1r14leoKAAgTWlqBZoeH-uYUAeizuQLkJec,53455
+brainstate/surrogate.py,sha256=zCNrJG-UMuvH7Le7iBuolDc-Oy8UY-xvw77tNhkGvK8,53861
 brainstate/transform.py,sha256=OvshYpPnp3YXPG6riy15Ve7jcX8t0aaATl1xZnsFeic,858
 brainstate/typing.py,sha256=cIlcgaWbvsFyEb5GgwZD7NHJh9qAp8l-N-rYN5W9FsU,10519
 brainstate/augment/__init__.py,sha256=Q9-JIwQ1FNn8VLS1MA9MrSylbvUjWSw98whrI3NIuKo,1229
-brainstate/augment/_autograd.py,sha256=zMMLZOidQq2p96wzEOgR-MynV2JH1l1AnoJ28eVwD-s,30100
-brainstate/augment/_autograd_test.py,sha256=UhNd41luca_Kj9a8byL3Dq64Ta55WU-je-LCiK_z0Vc,45060
-brainstate/augment/_eval_shape.py,sha256=mYRB5clSRHdXX0c4c5_WDIIBtRl4a9-xl8Tg60mgLBI,3854
+brainstate/augment/_autograd.py,sha256=XG0BgWrZ7LPelt2nRBKhpiMqz9jNS8NsK7x_I3DJINs,30884
+brainstate/augment/_autograd_test.py,sha256=1gR_38_i3pfH4S744G9x4KfGGtlyY0MBt43R_fndd9A,45396
+brainstate/augment/_eval_shape.py,sha256=jyzFzvZnUkbky7lxuchyFNGn0oJP1u10ZH5p79tAYMI,3868
 brainstate/augment/_eval_shape_test.py,sha256=WXrmZKmnykmYPveRfZbrDF0sFm2_PTr982yl4JM7Ebg,1390
-brainstate/augment/_mapping.py,sha256=AvGt8qhOdBgRJyWZn7J5JnOvxZc7XfZ_JKXaECOQjIw,43500
-brainstate/augment/_mapping_test.py,sha256=An5iSqrAXq_OVTvY3GFTjsB9AvFZ4j6AWNZQ_WWfra4,15319
+brainstate/augment/_mapping.py,sha256=R9CYp8JrdWpQcQE-XGr7POXyHqIEAYPxZfWqHF-wc70,43815
+brainstate/augment/_mapping_test.py,sha256=tEyXioA4_v4WGRFrpgUInyKyBjAVNLonWdHKVacmphU,21946
 brainstate/augment/_random.py,sha256=bkngsIk6Wwi3e46I7YSbBjLCGAz0Q3WuadUH4mqTjbY,5348
 brainstate/compile/__init__.py,sha256=fQtG316MLkeeu1Ssp54Kghw1PwbGK5gNq9yRVJu0wjA,1474
 brainstate/compile/_ad_checkpoint.py,sha256=HM6L90HU0N84S30uJpX8wqTO0ZcDnctqmwf8qFlkDYo,9358
 brainstate/compile/_ad_checkpoint_test.py,sha256=XGqC8kjwKd64RPgBWXHAE3t8NdIXVUqmGDfrg5HqoVw,1743
-brainstate/compile/_conditions.py,sha256=fFSJWRFZxriAZFASj2lvlXpxbH323XcfsoB9ni2CYZ0,10183
+brainstate/compile/_conditions.py,sha256=Pkvkc1Im7_ythKa9nwrUCNxYwkZri78I91MKbo1lEX8,10183
 brainstate/compile/_conditions_test.py,sha256=bh8_D49Fl-s_xFYkO4joZJPgFGhw_vd6Wo8QB1HDKwE,8640
 brainstate/compile/_error_if.py,sha256=l8flahWEui9-9IavtZQjyutIeR8DU0jieahRvAXUxuw,2651
 brainstate/compile/_error_if_test.py,sha256=PjTIQOflegqnqJb9kO5YqwjW1OfrBmAtWlY_xn8rtDY,1897
@@ -31,22 +31,22 @@ brainstate/compile/_loop_collect_return.py,sha256=qOYGoD2TMZd2Px6y241GWNwSjOLyaK
 brainstate/compile/_loop_collect_return_test.py,sha256=pEJdcOthEM17q5kXYhgR6JfzzqibijD_O1VPXVe_Ml4,1808
 brainstate/compile/_loop_no_collection.py,sha256=tPkSxee41VexWEALpN1uuT78BDrX3uT1FHv8WZtov4c,7549
 brainstate/compile/_loop_no_collection_test.py,sha256=ivavF59xep_g9bV1SSdXd5E1f6nhc7EUBfcbgHpxbfg,1419
-brainstate/compile/_make_jaxpr.py,sha256=fhaazC2Lj8-Gk3VpWLSytzVbP9dLkuTAn_1Y0SBtC9s,35864
+brainstate/compile/_make_jaxpr.py,sha256=3VYwfjmHSp5Uh_VEQVXYgxsUcpWcDN_2tNnZKh1dFV0,37736
 brainstate/compile/_make_jaxpr_test.py,sha256=xPusEJikMQRfoOmUFrugGcE5UpRm0giHoL_NPomN5rI,4791
-brainstate/compile/_progress_bar.py,sha256=iTp-CkcK5aMbiA0ZdlSfEXXrc4_anifYhyDzl19dksI,7488
+brainstate/compile/_progress_bar.py,sha256=AhAyyI_ckzgaj0PSj1ep1hq8rGzQLSyC2aVCYaT-e-o,7502
 brainstate/compile/_unvmap.py,sha256=9S42MeTmFJa8nfBI_AjEfrAdUsDmM3KFat59O8zXIEw,4120
 brainstate/compile/_util.py,sha256=QD7lvS4Zb3P2HPNIm6mG8Rl_Lk2tQj0znOCPcH95XnI,6275
 brainstate/functional/__init__.py,sha256=j6-3Er4fgqWpvntzYCZVB3e5hoz-Z3aqvapITCuDri0,1107
 brainstate/functional/_activations.py,sha256=3ym1oDdftVxYlHIXWdS7XONiDFaDklqokIioi4qrsHo,21688
-brainstate/functional/_activations_test.py,sha256=66PTMI99pfi1hr2rfFixg2OlUN0WME9628T7U2NF0iI,13475
+brainstate/functional/_activations_test.py,sha256=3YR43Jyls2wWYKI-phsf6b1ZdgZa578ob1TfkitroMI,13559
 brainstate/functional/_normalization.py,sha256=_F5u9iJm8KJfxlnmV_LwJXtrAZOJj357aXc_sfGe4CQ,2567
 brainstate/functional/_others.py,sha256=iAqvLCmXQz6C80EjE0-5bwXCBKHoKqsggneiOAXNmQo,1699
 brainstate/functional/_spikes.py,sha256=YJ4pa5LlywximSE6ZAAs9g2k-8cO95mwxdD29HhT-oM,7051
 brainstate/graph/__init__.py,sha256=noo4TjBg6iEhjjwk0sAGUhR7Ge-z8Vnc2rLYUvnqttw,1295
 brainstate/graph/_graph_node.py,sha256=kl0CDKIcbXd9YzrZIKDhDeT8bboaoR5_OPwtL90eUEo,6835
 brainstate/graph/_graph_node_test.py,sha256=g6nrCCI-awfxhZLnjmYhfQvynumWJjz7fPxgIF-VSr4,2832
-brainstate/graph/_graph_operation.py,sha256=h2rdydWr9AHBMVlyYpBN9GdbVkPIgyjsctZXa8EmxQI,63401
-brainstate/graph/_graph_operation_test.py,sha256=nxMwEoaRZpTkIeR3h_Aduuys0Bm1DRlYLd_CXkErTSg,19715
+brainstate/graph/_graph_operation.py,sha256=nIOW0GsA3Yeh8X_ijw5wx4hQgY2sh4knVMt5ouG2Sn0,63975
+brainstate/graph/_graph_operation_test.py,sha256=xgHoHXAafEhOqISDKLqYxb-wmMrUpv2PLtH5OLZausI,19729
 brainstate/init/__init__.py,sha256=R1dHgub47o-WJM9QkFLc7x_Q7GsyaKKDtrRHTFPpC5g,1097
 brainstate/init/_base.py,sha256=A5VIJV_xPyCD4Lo4RXe2WPlWvADAy_UPhexuwzp-6EI,1633
 brainstate/init/_generic.py,sha256=o60EPJZe7laJgPsgrnorl3Km_Lc_Lr4qQQKmu6g5G4o,7992
@@ -55,36 +55,36 @@ brainstate/init/_random_inits_test.py,sha256=c1pUdWXQxL-seEvjO_8fBkaDL6T63b6ALNf
 brainstate/init/_regular_inits.py,sha256=tcZDTXR8jiRqERuJleg_9JRGGzJj-V3LhTMcr_gt_k0,3149
 brainstate/init/_regular_inits_test.py,sha256=bvv0AOfLOEP0BIQIBLztKw3EPyEp7n2fHW8PUrmWuHk,1820
 brainstate/nn/__init__.py,sha256=nDIBDlzHaBy-vk4cb7FiPO6zoNe2tI5qvHIK8O7yrOU,3721
-brainstate/nn/_collective_ops.py,sha256=rWcjqaP0rW6sXdhE0fjtDi4twmq5VvhNEg8AoiQ-tDU,17390
+brainstate/nn/_collective_ops.py,sha256=Mmg-0ImuGvpv2qggu5wfBrA0MLBaq96D0P5DCZ3yolY,17418
 brainstate/nn/_collective_ops_test.py,sha256=bwq0DApcsk0_2xpxMl0_e2cGKT63g5rSngpigCm07ps,1409
-brainstate/nn/_common.py,sha256=Pt8P1LgE0qW3QnfX8CQQGH3yZ78RL7NkIqqYEw_z8xs,7096
+brainstate/nn/_common.py,sha256=qHAOID_eeKiPUXk_ION65sbYXyF-ddH5w5BayvH8Thg,6431
 brainstate/nn/_conv.py,sha256=Zk-yj34n6CkjntcM9xpMGLTxKNfWdIWsTsoGbtdL0yU,18448
 brainstate/nn/_conv_test.py,sha256=2lcUTG7twkyhuyKwuBux-NgU8NU_W4Cp1-G8EyDJ_uk,8862
 brainstate/nn/_delay.py,sha256=l36FBgNhfL64tM3VGOsJNTtKr44HjxxtBWMFFCm3Pks,17361
 brainstate/nn/_dropout.py,sha256=Dq3hQrOBT6gODlDbcoag6zLGXTU_p_2MhnmVsV57Hds,17783
-brainstate/nn/_dropout_test.py,sha256=9i2ZW5by0S9zfVHmZKLN0WpI1bXH6h9-QLXamG6lVXA,4432
-brainstate/nn/_dynamics.py,sha256=Rx4Hwl87H94DgMM644hDi3nN9D0wvXkpWrog_Ek1sr4,47734
+brainstate/nn/_dropout_test.py,sha256=L46PvC2OA7EnS4MsRhh_YnvKheHYNafOsKM8uzux_zo,4446
+brainstate/nn/_dynamics.py,sha256=PTTGDqe-uRvrTZbhNGy1O4hT_wG4vxbOXyfbOx2itaw,48341
 brainstate/nn/_dynamics_test.py,sha256=w7AV57LdhbBNYprdFpKq8MFSCbXKVkGgp_NbL3ANX3I,2769
-brainstate/nn/_elementwise.py,sha256=4czeJWGQopV49iZo8DuN_WzAbXoMC1gtqaGjlON6e7c,33291
+brainstate/nn/_elementwise.py,sha256=l3Wk65UDLFc3xFH1jwTXRwKAVmt3dUeG0iF_Al0w1_0,33487
 brainstate/nn/_elementwise_test.py,sha256=_dd9eX2ZJ7p24ahuoapCaRTZ0g1boufXMyqHFx1d4WY,5688
 brainstate/nn/_embedding.py,sha256=SaAJbgXmuJ8XlCOX9ob4yvmgh9Fk627wMguRzJMJ1H8,2138
-brainstate/nn/_exp_euler.py,sha256=ndDB43PM4jsZKu_zdLTZ2-ojnuNrg55LZap23oBTtdA,3493
+brainstate/nn/_exp_euler.py,sha256=WTpZm-XQmsdMLNazY7wIu8eeO6pK0kRzt2lJnhEgMIk,3293
 brainstate/nn/_exp_euler_test.py,sha256=XD--qMbGHrHa3WtcPMmJKk59giDcEhSqZuBOmTNYUr8,1227
 brainstate/nn/_fixedprob.py,sha256=KGXohiU0wZnFIQDuwiRUTFsbsr8R0p8zgi5UZDuv1Bk,10004
 brainstate/nn/_fixedprob_test.py,sha256=qbRBh-MpMtEOsg492gFu2w9-FOP9z_bXapm-Q0gLLYM,3929
-brainstate/nn/_inputs.py,sha256=wPOfPE4IesNoDmxZJxqR0siBlJioEX-_1IZ2cltAIpM,20605
-brainstate/nn/_linear.py,sha256=5WuhcqU-uBUC91vnwezQYMHPKmlZPDgIJ5UpffxoX1I,14472
+brainstate/nn/_inputs.py,sha256=hMPkx9qDBpJWPshZXLF4H1QiYK1-46wntHUIlG7cT7c,20603
+brainstate/nn/_linear.py,sha256=FnPxATdT66DecjTW0tUTyL6clQmN_cG8kPC3qDWUE6A,14500
 brainstate/nn/_linear_mv.py,sha256=6hDXx4yPqRSa7uIsW9f9eJuy23dcXN9Mp2_lSvw8BDA,2635
 brainstate/nn/_linear_mv_test.py,sha256=ZCM1Zy6mImQfCfdZOGnTwkiLLPXK5yalv1Ts9sWZuPA,3864
 brainstate/nn/_linear_test.py,sha256=eIS-VCR3QmXB_byO1Uexg65Pv48CBRUA_Je-UGrFVTY,2925
 brainstate/nn/_ltp.py,sha256=_najNUyfaFYcOlUTm7ThJopInbos3kwJyrm-QUfI-hc,861
-brainstate/nn/_module.py,sha256=tleUiEd3ytEx4fZLu9NButELVlbKZNIQbXX78bdK8QU,12770
-brainstate/nn/_module_test.py,sha256=_RKHY8LHg6nbFHkvPAsEPVpp9rW6XfrcF8YD68u4o4w,8956
+brainstate/nn/_module.py,sha256=jlFaoltT2F8a7cxsyu6fXp7mfIkEsk-JD3G1Xh3Ay8I,12783
+brainstate/nn/_module_test.py,sha256=sjT7t-N4ZZKPN_MujNv8bcT5uZQD-CmU5nKz9KwlBUc,8963
 brainstate/nn/_neuron.py,sha256=2walTScvL034LS53pArDASXz6z26SSPbmCvchWWjkUU,27441
 brainstate/nn/_neuron_test.py,sha256=QF8pixUqA5Oj7MrNi2NR8VAnfGpAvNpwV2mBc3e_pTY,6393
-brainstate/nn/_normalizations.py,sha256=4B1YVI6AZyJLHseJK9eiSeMOIq--0ZuDG58iC9mrekE,37387
-brainstate/nn/_normalizations_test.py,sha256=JTqpH265GT0eKw3hXmT0qC8ZM5NkUe98of0uYS7_5Us,2456
-brainstate/nn/_poolings.py,sha256=h5aIj_K3X2HSMjrOtkgwiCpvLxfXXEUvFFa57Cxp5bI,46969
+brainstate/nn/_normalizations.py,sha256=YSC1W7JaexoZ8tKcy5B-dm-_x8GuPzS_6XBfkaKpdXM,37464
+brainstate/nn/_normalizations_test.py,sha256=56YfszPQW9y3VoX90-seyLeWDX8akYQTRLNBvrobmS0,2498
+brainstate/nn/_poolings.py,sha256=yKG9BLnvDvZb7HMrAt4Y1sLCjFQ8UmWfSSQmICAPWag,47137
 brainstate/nn/_poolings_test.py,sha256=qje9PVWvPGiYOv6UlTEWfpqqjpv4Xop5rO0ATcgcF0w,7497
 brainstate/nn/_projection.py,sha256=_TW-YmfPqUDkekMMK7tXrdLbjT4Ou2sUAecTNNn7FVM,17498
 brainstate/nn/_rate_rnns.py,sha256=OQHMyq9kKd2q44dOHLuTqHKTGuraPMqXVutFP_LYIyU,20676
@@ -92,40 +92,40 @@ brainstate/nn/_rate_rnns_test.py,sha256=__hhx7e6LX_1mDLLQyIi4TNCaFAWnOVSTIgwHNjz
 brainstate/nn/_readout.py,sha256=OJjSba5Wr7dtUXqYhAv1D7BUGOI-lAmg6urxPBrZe3c,7116
 brainstate/nn/_readout_test.py,sha256=L2T0-SkiACxkY_I5Pbnbmy0Zw3tbpV3l5xVzAw42f2g,2136
 brainstate/nn/_stp.py,sha256=-ahDEqSp8bQsU_nUK4jks8fjMYKgIbO0v7zpyGVuXtA,8645
-brainstate/nn/_synapse.py,sha256=liKhmHZFQpEdu1R76AljeJMAKz5HK8Yq95ENa6FLVAc,20145
+brainstate/nn/_synapse.py,sha256=E43DxH0PW3pa6rKn9UWkD_PxIdpMAo04Yy41gaR3KPc,19868
 brainstate/nn/_synapse_test.py,sha256=xmCWFxZUIM2YtmW5otKnADGCCK__4JpXmSYcZ3wzlQM,4994
 brainstate/nn/_synaptic_projection.py,sha256=UFgzsMB1VZ9ieumwmaTIC7irLrZ4pmwfiuOYomkwXG4,17917
 brainstate/nn/_synouts.py,sha256=jWQP1-qXFpdYgyUSJNFD7_bk4_-67ok36br-OzbcSXY,4524
 brainstate/nn/_synouts_test.py,sha256=sfjotlS--4hT22Vb5RfmpfmXABa8z-VdxZYUSocMmlU,2244
-brainstate/nn/_utils.py,sha256=epfELIy1COgdS9z5be-fmbFhagNugcIHpw4ww-HlkSY,3123
-brainstate/nn/metrics.py,sha256=TXCB_yGQzamklJCI5FGOAZ5dihtY-gjYlSi2SCLC3LA,14700
+brainstate/nn/_utils.py,sha256=nmp6MRQoKJhfxkbbnZJEyhyMQ4DZh8FNN_JQsl55kHE,3130
+brainstate/nn/metrics.py,sha256=kFhUZRowtk4IoBIOPbENO4KezA6ZW_AxLAVknKFiL-s,14728
 brainstate/optim/__init__.py,sha256=7Ao0LCtDNAoxSRSXiLLKnd1_4mR2GSExizpN38il-Fo,1195
 brainstate/optim/_base.py,sha256=P37k8w46iQZZZnFLa5OF83Sb8DLQfYMez9ZRObBqVsE,1835
 brainstate/optim/_lr_scheduler.py,sha256=0t7kl35MLmc0UcSu587nI28cI3SohYhy_19Al-5dNTM,15289
 brainstate/optim/_lr_scheduler_test.py,sha256=gzzd_rzgtDWzF-93mujg043jv2wO9IceXGKiIYw5_ko,1790
-brainstate/optim/_optax_optimizer.py,sha256=RcO5N5JO_U0P4Q82waQuwMS7yWR-oP18eUwgzXJX0EY,5274
+brainstate/optim/_optax_optimizer.py,sha256=eNaHge-9URe82WvIGDGrWrAr0MAkYWe4eCsTwijJlbU,5337
 brainstate/optim/_optax_optimizer_test.py,sha256=GVcszbNxBZ4z54bS5r-zDzP8gQZl7s5IsEZEd6G2gLo,1775
 brainstate/optim/_sgd_optimizer.py,sha256=Fn4CdGaDfV4qc_4gacCcGSjit5-ZrsBbQP-KB0Ocdac,46110
 brainstate/random/__init__.py,sha256=c5q-RC3grRIjx-HBb2IhKZpi_xzbFmUUxzRAzqfREic,1045
-brainstate/random/_rand_funs.py,sha256=c4xiY2NeMizSslxbWOa-QJZ3h-LDfsgH4fbvBMSNLWg,137646
-brainstate/random/_rand_funs_test.py,sha256=G8BuxDjBSeE-Mh7KuwVk6mPQhHE0m1R5HDFljEOIdzg,20669
-brainstate/random/_rand_seed.py,sha256=1ZdfFZWyOhpd72EDdEDmpkp3yoLVwdv-sGI9BwiZfzI,5949
+brainstate/random/_rand_funs.py,sha256=0OY5kqEVrA2im-uKW2Rown0iIy7F0u54hGXHKxAP-go,137905
+brainstate/random/_rand_funs_test.py,sha256=1K6bT6hhOPhomE-AMyPpT6MtG6ewANRSb5bJSU-ysYc,20690
+brainstate/random/_rand_seed.py,sha256=ROpPgopGyL2ITNWJAM25NYjTyBizUoC4LzPUX-JpTpk,5998
 brainstate/random/_rand_seed_test.py,sha256=waXXfch57X1XE1zDnCRokT6ziZOK0g-lYE80o6epDYM,1536
-brainstate/random/_rand_state.py,sha256=dTO7wqmuTYRdPy7ItsrK-7aNt5QQXOsZ4XwiH7mzmy8,55170
+brainstate/random/_rand_state.py,sha256=hQ_govRxfIgsmtFe_V2B-jftqiMWRWrd9wkviPFOziY,55523
 brainstate/random/_random_for_unit.py,sha256=kGp4EUX19MXJ9Govoivbg8N0bddqOldKEI2h_TbdONY,2057
 brainstate/util/__init__.py,sha256=6efwr63osmqviNU_6_Nufag19PwxRDFvDAZrq6sH5yo,1555
-brainstate/util/_pretty_pytree_test.py,sha256=Dn0TdjX6wLBXaTD4jfYTu6cKfFHwKSxi4_3bX7kB_IA,5621
 brainstate/util/caller.py,sha256=RBRwu1HT_Lww7EE3WTjzrwwKqVWg7dn_JXKmO9ojM38,2752
 brainstate/util/error.py,sha256=GBLaU_glQlgdrjBYIMiO3Dy5N_ST6-QNsgAOa9kjUjw,1757
 brainstate/util/filter.py,sha256=skuq51y0P5-N611fQ-zB2QyaJIFqN_WbROzZHv0RgOc,14086
 brainstate/util/others.py,sha256=vGo4uW0vd7wL5DYtxKWqXhoNmI-O_p26HsucoMnhdVA,16407
-brainstate/util/pretty_pytree.py,sha256=HmXZpKGps1Rs0DEN7r1eYN6xMM6FgDpCC8SdEHjOsvo,33534
+brainstate/util/pretty_pytree.py,sha256=ROrOO03fxq2NvsZHlz1XG7B4NhAjHXavPjh179GUwG4,33604
+brainstate/util/pretty_pytree_test.py,sha256=64x2c9teehF_UoozCuuJ9WQj_287cTOTV1x58jxOTv8,5914
 brainstate/util/pretty_repr.py,sha256=7Xp7IFNUeP7cGlpvwwJyBslbQVnXEqC1I6neV1Jx1S8,10527
 brainstate/util/pretty_table.py,sha256=uJVaamFGQ4nKP8TkEGPWXHpzjMecDo2q1Ah6XtRjdPY,108117
 brainstate/util/scaling.py,sha256=U6DM-afPrLejiGqo1Nla7z4YbTBVicctsBEweurr_mk,7524
-brainstate/util/struct.py,sha256=7HbbQNrZ3zxYw93MU1bUZ9ZPBKftYOVKuXEochLSErw,17479
-brainstate-0.1.4.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.1.4.dist-info/METADATA,sha256=EWarKpuYFIN0ITFn8lwsUy0h9-KjrVPDwcvjQMHFRto,4135
-brainstate-0.1.4.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.1.4.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.1.4.dist-info/RECORD,,
+brainstate/util/struct.py,sha256=2Y_wuDFQ6ldl_H4_w0IjzAtkbHooVgdsVbnT7Z6_Efc,17528
+brainstate-0.1.6.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.1.6.dist-info/METADATA,sha256=PsI82NGU6Zeh18JitOiA0jnEgullQ5aA_tl4zdl4Iqw,4122
+brainstate-0.1.6.dist-info/WHEEL,sha256=AHX6tWk3qWuce7vKLrj7lnulVHEdWoltgauo8bgCXgU,109
+brainstate-0.1.6.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.1.6.dist-info/RECORD,,

{brainstate-0.1.4.dist-info → brainstate-0.1.6.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.38.4)
+Generator: setuptools (75.1.0)
 Root-Is-Purelib: true
 Tag: py2-none-any
 Tag: py3-none-any