brainstate 0.1.0.post20250120__py2.py3-none-any.whl → 0.1.0.post20250127__py2.py3-none-any.whl

brainstate/__init__.py

@@ -22,7 +22,6 @@ __version__ = "0.1.0"
 from . import augment
 from . import compile
 from . import environ
-from . import event
 from . import functional
 from . import graph
 from . import init
@@ -39,7 +38,7 @@ from ._state import __all__ as _state_all
 
 __all__ = (
     [
-        'augment', 'compile', 'environ', 'event', 'functional',
+        'augment', 'compile', 'environ', 'functional',
         'graph', 'init', 'mixin', 'nn', 'optim', 'random',
         'surrogate', 'typing', 'util',
         # deprecated

brainstate/augment/__init__.py

@@ -17,24 +17,14 @@
 This module includes transformations for augmenting the functionalities of JAX code.
 """
 
-from ._autograd import *
-from ._autograd import __all__ as _autograd_all
-from ._eval_shape import *
-from ._eval_shape import __all__ as _eval_shape_all
-from ._mapping import *
-from ._mapping import __all__ as _mapping_all
-from ._random import *
-from ._random import __all__ as _random_all
+from ._autograd import GradientTransform, grad, vector_grad, hessian, jacobian, jacrev, jacfwd
+from ._eval_shape import abstract_init
+from ._mapping import vmap, pmap, map
+from ._random import restore_rngs
 
-__all__ = (
-    _eval_shape_all
-    + _autograd_all
-    + _mapping_all
-    + _random_all
-)
-del (
-    _eval_shape_all,
-    _autograd_all,
-    _mapping_all,
-    _random_all
-)
+__all__ = [
+    'GradientTransform', 'grad', 'vector_grad', 'hessian', 'jacobian', 'jacrev', 'jacfwd',
+    'abstract_init',
+    'vmap', 'pmap', 'map',
+    'restore_rngs',
+]

brainstate/compile/__init__.py

@@ -17,41 +17,22 @@
 This module contains the functions for the compilation of JAX code.
 """
 
-from ._ad_checkpoint import *
-from ._ad_checkpoint import __all__ as _ad_checkpoint_all
-from ._conditions import *
-from ._conditions import __all__ as _conditions_all
-from ._error_if import *
-from ._error_if import __all__ as _jit_error_all
-from ._jit import *
-from ._jit import __all__ as _jit_all
-from ._loop_collect_return import *
-from ._loop_collect_return import __all__ as _loops_collection
-from ._loop_no_collection import *
-from ._loop_no_collection import __all__ as _loops_no_collection
-from ._make_jaxpr import *
-from ._make_jaxpr import __all__ as _make_jaxpr_all
-from ._progress_bar import *
-from ._progress_bar import __all__ as _progress_bar_all
+from ._ad_checkpoint import checkpoint, remat
+from ._conditions import cond, switch, ifelse
+from ._error_if import jit_error_if
+from ._jit import jit
+from ._loop_collect_return import scan, checkpointed_scan, for_loop, checkpointed_for_loop
+from ._loop_no_collection import while_loop, bounded_while_loop
+from ._make_jaxpr import StatefulFunction, make_jaxpr
+from ._progress_bar import ProgressBar
 
-__all__ = (
-    _jit_error_all
-    + _conditions_all
-    + _make_jaxpr_all
-    + _jit_all
-    + _progress_bar_all
-    + _loops_collection
-    + _loops_no_collection
-    + _ad_checkpoint_all
-)
-
-del (
-    _jit_error_all,
-    _conditions_all,
-    _loops_collection,
-    _make_jaxpr_all,
-    _jit_all,
-    _progress_bar_all,
-    _loops_no_collection,
-    _ad_checkpoint_all
-)
+__all__ = [
+    'checkpoint', 'remat',
+    'cond', 'switch', 'ifelse',
+    'jit_error_if',
+    'jit',
+    'scan', 'checkpointed_scan', 'for_loop', 'checkpointed_for_loop',
+    'while_loop', 'bounded_while_loop',
+    'StatefulFunction', 'make_jaxpr',
+    'ProgressBar',
+]

brainstate/compile/_make_jaxpr.py

@@ -73,6 +73,13 @@ from brainstate._state import State, StateTraceStack
 from brainstate._utils import set_module_as
 from brainstate.typing import PyTree
 
+
+if jax.__version_info__ < (0, 4, 38):
+    from jax.core import ClosedJaxpr
+else:
+    from jax.extend.core import ClosedJaxpr
+
+
 AxisName = Hashable
 
 __all__ = [
@@ -660,7 +667,7 @@ def _make_jaxpr(
     axis_env: Sequence[tuple[AxisName, int]] | None = None,
     return_shape: bool = False,
     abstracted_axes: Any | None = None,
-) -> Callable[..., (jax.core.ClosedJaxpr | tuple[jax.core.ClosedJaxpr, Any])]:
+) -> Callable[..., (ClosedJaxpr | tuple[ClosedJaxpr, Any])]:
     """Creates a function that produces its jaxpr given example args.
 
     Args:
@@ -723,7 +730,7 @@ def _make_jaxpr(
     def _abstractify(args, kwargs):
         flat_args, in_tree = jax.tree.flatten((args, kwargs))
         if abstracted_axes is None:
-            return map(jax.api_util.shaped_abstractify, flat_args), in_tree, [True] * len(flat_args)
+            return map(shaped_abstractify, flat_args), in_tree, [True] * len(flat_args)
         else:
             axes_specs = _flat_axes_specs(abstracted_axes, *args, **kwargs)
             in_type = pe.infer_lambda_input_type(axes_specs, flat_args)

brainstate/compile/_make_jaxpr_test.py

@@ -18,12 +18,16 @@ from __future__ import annotations
 import unittest
 
 import jax
-import jax.extend as je
 import jax.numpy as jnp
 import pytest
 
 import brainstate as bst
 
+if jax.__version_info__ < (0, 4, 38):
+    from jax.core import jaxpr_as_fun
+else:
+    from jax.extend.core import jaxpr_as_fun
+
 
 class TestMakeJaxpr(unittest.TestCase):
     def test_compar_jax_make_jaxpr(self):
@@ -85,7 +89,7 @@ class TestMakeJaxpr(unittest.TestCase):
         print(jaxpr)
         jaxpr, _ = bst.compile.make_jaxpr(f3)(jnp.zeros(1))
         print(jaxpr)
-        self.assertTrue(jnp.allclose(je.core.jaxpr_as_fun(jaxpr)(jnp.zeros(1), st1.value)[0],
+        self.assertTrue(jnp.allclose(jaxpr_as_fun(jaxpr)(jnp.zeros(1), st1.value)[0],
                                      f3(jnp.zeros(1))))
 
     def test_compar_jax_make_jaxpr2(self):
@@ -103,10 +107,10 @@ class TestMakeJaxpr(unittest.TestCase):
         print()
         print(jaxpr)
         print(states)
-        print(jax.core.jaxpr_as_fun(jaxpr)(jnp.zeros(1), st1.value))
+        print(jaxpr_as_fun(jaxpr)(jnp.zeros(1), st1.value))
         jaxpr = jax.make_jaxpr(ffa)(jnp.zeros(1))
         print(jaxpr)
-        print(jax.core.jaxpr_as_fun(jaxpr)(jnp.zeros(1)))
+        print(jaxpr_as_fun(jaxpr)(jnp.zeros(1)))
 
     def test_compar_jax_make_jaxpr3(self):
         def fa(x):
@@ -116,10 +120,10 @@ class TestMakeJaxpr(unittest.TestCase):
         print()
         print(jaxpr)
         print(states)
-        # print(jax.core.jaxpr_as_fun(jaxpr)(jnp.zeros(1)))
+        # print(jaxpr_as_fun(jaxpr)(jnp.zeros(1)))
         jaxpr = jax.make_jaxpr(fa)(jnp.zeros(1))
         print(jaxpr)
-        # print(jax.core.jaxpr_as_fun(jaxpr)(jnp.zeros(1)))
+        # print(jaxpr_as_fun(jaxpr)(jnp.zeros(1)))
 
 
 def test_return_states():

brainstate/compile/_progress_bar.py

@@ -35,6 +35,47 @@ Output = Any
 class ProgressBar(object):
     """
     A progress bar for tracking the progress of a jitted for-loop computation.
+
+    It can be used in :py:func:`for_loop`, :py:func:`checkpointed_for_loop`, :py:func:`scan`,
+    and :py:func:`checkpointed_scan` functions. Or any other jitted function that uses
+    a for-loop.
+
+    The message displayed in the progress bar can be customized by the following two methods:
+
+    1. By passing a string to the `desc` argument. For example:
+
+    .. code-block:: python
+
+            ProgressBar(desc="Running 1000 iterations")
+
+    2. By passing a tuple with a string and a callable function to the `desc` argument. The callable
+       function should take a dictionary as input and return a dictionary. The returned dictionary
+       will be used to format the string. For example:
+
+    .. code-block:: python
+
+                a = bst.State(1.)
+                def loop_fn(x):
+                    a.value = x.value + 1.
+                    return jnp.sum(x ** 2)
+
+                pbar = ProgressBar(desc=("Running {i} iterations, loss = {loss}",
+                                         lambda i_carray_y: {"i": i_carray_y["i"], "loss": i_carray_y["y"]}))
+
+                bst.compile.for_loop(loop_fn, xs, pbar=pbar)
+
+    In this example, ``"i"`` denotes the iteration number and ``"loss"`` is computed from the output,
+    the ``"carry"`` is the dynamic state in the loop, for example ``a.value`` in this case.
+
+
+    Args:
+        freq: The frequency at which to print the progress bar. If not specified, the progress
+            bar will be printed every 5% of the total iterations.
+        count: The number of times to print the progress bar. If not specified, the progress
+            bar will be printed every 5% of the total iterations.
+        desc: A description of the progress bar. If not specified, a default message will be
+            displayed.
+        kwargs: Additional keyword arguments to pass to the progress bar.
     """
     __module__ = "brainstate.compile"
 
@@ -42,7 +83,7 @@ class ProgressBar(object):
         self,
         freq: Optional[int] = None,
         count: Optional[int] = None,
-        desc: Optional[Tuple[str, Callable[[Dict], Dict]]] = None,
+        desc: Optional[Tuple[str, Callable[[Dict], Dict]] | str] = None,
         **kwargs
     ):
         # print rate
@@ -62,9 +103,12 @@ class ProgressBar(object):
 
         # description
         if desc is not None:
-            assert isinstance(desc, (tuple, list)), 'Description should be a tuple or list.'
-            assert isinstance(desc[0], str), 'Description should be a string.'
-            assert callable(desc[1]), 'Description should be a callable.'
+            if isinstance(desc, str):
+                pass
+            else:
+                assert isinstance(desc, (tuple, list)), 'Description should be a tuple or list.'
+                assert isinstance(desc[0], str), 'Description should be a string.'
+                assert callable(desc[1]), 'Description should be a callable.'
         self.desc = desc
 
         # check if tqdm is installed
@@ -136,8 +180,7 @@ class ProgressBarRunner(object):
         self.tqdm_bars[0].close()
 
     def __call__(self, iter_num, **kwargs):
-        data = dict(i=iter_num, **kwargs)
-        data = dict() if isinstance(self.message, str) else self.message[1](data)
+        data = dict() if isinstance(self.message, str) else self.message[1](dict(i=iter_num, **kwargs))
         assert isinstance(data, dict), 'Description function should return a dictionary.'
 
         _ = jax.lax.cond(

brainstate/compile/_unvmap.py

@@ -78,7 +78,7 @@ mlir.register_lowering(
 
 # unvmap_any
 
-unvmap_any_p = jax.core.Primitive("unvmap_any")
+unvmap_any_p = Primitive("unvmap_any")
 
 
 def unvmap_any(x):
@@ -109,7 +109,7 @@ mlir.register_lowering(
 
 # unvmap_max
 
-unvmap_max_p = jax.core.Primitive("unvmap_max")
+unvmap_max_p = Primitive("unvmap_max")
 
 
 def unvmap_max(x):
@@ -156,7 +156,7 @@ def _without_vmap_batch(x, batch_axes):
     return _without_vmap(x), batching.not_mapped
 
 
-_no_vmap_prim = jax.core.Primitive('no_vmap')
+_no_vmap_prim = Primitive('no_vmap')
 _no_vmap_prim.def_impl(_without_vmap_imp)
 _no_vmap_prim.def_abstract_eval(_without_vmap_abs)
 batching.primitive_batchers[_no_vmap_prim] = _without_vmap_batch

brainstate/graph/__init__.py

@@ -14,16 +14,16 @@
 # ==============================================================================
 
 
-from ._graph_node import *
-from ._graph_node import __all__ as _graph_node__all__
-from ._graph_operation import *
-from ._graph_operation import __all__ as _graph_operation__all__
-
-__all__ = (
-    _graph_node__all__ +
-    _graph_operation__all__
-)
-del (
-    _graph_node__all__,
-    _graph_operation__all__
+from ._graph_node import Node, Dict, List, Sequential
+from ._graph_operation import (
+    pop_states, nodes, states, treefy_states, update_states, flatten, unflatten,
+    treefy_split, treefy_merge, iter_leaf, iter_node, clone, graphdef,
+    call, RefMap, GraphDef, NodeRef, NodeDef
 )
+
+__all__ = [
+    'Node', 'Dict', 'List', 'Sequential',
+    'pop_states', 'nodes', 'states', 'treefy_states', 'update_states', 'flatten', 'unflatten',
+    'treefy_split', 'treefy_merge', 'iter_leaf', 'iter_node', 'clone', 'graphdef',
+    'call', 'RefMap', 'GraphDef', 'NodeRef', 'NodeDef',
+]

brainstate/nn/_dyn_impl/_inputs.py

@@ -224,7 +224,8 @@ def poisson_input(
     weight = maybe_state(weight)
 
     assert isinstance(target, State), 'The target must be a State.'
-    p = (freq * environ.get_dt()).to_decimal()
+    p = freq * environ.get_dt()
+    p = p.to_decimal() if isinstance(p, u.Quantity) else p
     a = num_input * p
     b = num_input * (1 - p)
     tar_val = target.value
@@ -291,7 +292,8 @@ def poisson_input(
 
         # update target variable
         target.value = jax.tree.map(
-            lambda x: x * weight,
+            lambda tar, x: tar + x * weight,
+            target.value,
             inp,
             is_leaf=u.math.is_quantity
         )

brainstate/nn/_elementwise/_dropout_test.py

@@ -28,7 +28,7 @@ class TestDropout(unittest.TestCase):
         dropout_layer = bst.nn.Dropout(0.5)
 
         # Input data
-        input_data = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
+        input_data = np.arange(20)
 
         with bst.environ.context(fit=True):
             # Apply dropout

{brainstate-0.1.0.post20250120.dist-info → brainstate-0.1.0.post20250127.dist-info}/METADATA

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

@@ -1,4 +1,4 @@
-brainstate/__init__.py,sha256=A-QKdOvSalsCMxgk80Iz6_xMiUin6con6JaONHfciSY,1526
+brainstate/__init__.py,sha256=AkZyyFkn4fB8g2aT6Rc2MO1xICPpUZuDtdze-eUQNc0,1496
 brainstate/_state.py,sha256=GZ46liHZSHbAHQEuELvOeoJ27P9xiZDz06G2AASjAjA,29142
 brainstate/_state_test.py,sha256=rJUFRSXEqrrl4qANRewY9mnDlzSbtHwBIGeZ0ku-8Dg,1650
 brainstate/_utils.py,sha256=uJ6WWKq3yb05ZdktCQGLWOXsOJveL1H9pR7eev70Jes,1693
@@ -9,7 +9,7 @@ brainstate/mixin_test.py,sha256=Oq_0fwC9vpXDN4t4dTBhWzLdFDNlcYsrcip14F1yECI,3079
 brainstate/surrogate.py,sha256=t4SzVwUVMAPtC-O1vFbuE9F4265wgAv7ud77ufIJsuk,48464
 brainstate/transform.py,sha256=cxbymTlJ6uHvJWEEYXzFUkAySs_TbUTHakt0NQgWJ3s,808
 brainstate/typing.py,sha256=Qh-LBzm6oG4rSXv4V5qB8SNYcoOR7bASoK_iQxnlafk,10467
-brainstate/augment/__init__.py,sha256=BtXIBel7GbttmfBX6grxOxl0IiOJxLEa7qCGAXumamE,1286
+brainstate/augment/__init__.py,sha256=zGPq1eTB_56GRCNC9TiPLKTw07PA2O0OCi7bgjYIrY4,1193
 brainstate/augment/_autograd.py,sha256=o9ivoEY7BmtdM1XmzdMmeRXpj6Tvn5xNB8LSGp2HKC8,25238
 brainstate/augment/_autograd_test.py,sha256=S2eEgrwTzdSi3u2nKE3u37WSThosLwx1WCP9ptJAGKo,44060
 brainstate/augment/_eval_shape.py,sha256=ObCgsZ704kLduB1dbjJZh5nVQYEkLR5ebK74V5NV42k,3892
@@ -17,7 +17,7 @@ brainstate/augment/_eval_shape_test.py,sha256=LFOJx7CWltmRLXdGY175UebLwtEMz2CzJ_
 brainstate/augment/_mapping.py,sha256=nU6Y7fSnYXyQSILXU2QT-O73Fm3pnwOmgUoDaHqjve8,21544
 brainstate/augment/_mapping_test.py,sha256=_KFhE3CXItwpbZ1gJfrDu3yUtX0YbfPUuHJG_G_BXEs,8963
 brainstate/augment/_random.py,sha256=rkB4w4BkKsz9p8lTk31kVHvlVPJSvtGk8REn936KI_4,3071
-brainstate/compile/__init__.py,sha256=qZZIYoyEl51IFkFu-Hb-bP3PAEHo94HlTDf57P2ze08,1858
+brainstate/compile/__init__.py,sha256=fQtG316MLkeeu1Ssp54Kghw1PwbGK5gNq9yRVJu0wjA,1474
 brainstate/compile/_ad_checkpoint.py,sha256=K6I4vnznDsqC9cUeCnez9UdV9r_toGA3zHezoHLA6mI,9377
 brainstate/compile/_ad_checkpoint_test.py,sha256=R1I76nG4zIqb6g3M_VxWts7rUC1OHJCjtQhPkcbXodk,1746
 brainstate/compile/_conditions.py,sha256=gApsHKGQrf1QBjoKXDVL7VsoeJ2zFtSc-hFz9nbYcF0,10113
@@ -30,32 +30,18 @@ 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=DQf_80w3p0wi2Gb9P6_tLMJ0Oadgyr_jWkVjus0MSjw,33205
-brainstate/compile/_make_jaxpr_test.py,sha256=3XaX8LUuG6UjolcD83qDVo5odf8FCDppdr9Q6V0NBs4,4303
-brainstate/compile/_progress_bar.py,sha256=0oVlZ4kW_ZMciJjOR_ebj3PNe_XkCMkoQpv-HUUdoF0,5554
-brainstate/compile/_unvmap.py,sha256=EY4rbqCzzPOiaRwpWTiyBwb5dVkYFnacHhBZUZObxPI,4255
+brainstate/compile/_make_jaxpr.py,sha256=J4oWoPBwG-fdJvNhBEtNgmo3rXrIWCoajELhaIumgPU,33309
+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=aCvkTV--g4NsqcodTdBAISt4EwgezCbKzNUV58n-Q_Y,6304
-brainstate/event/__init__.py,sha256=gSEem-1oTHgy99Mjm3uumTXVd93tLVl0c4dUgRpoifk,895
-brainstate/event/_csr.py,sha256=PYKw8CGNgQ24MxQDoeBZTrPuC7Z-GetXQld9KiTbNYw,40063
-brainstate/event/_csr_benchmark.py,sha256=xrj2DSWzw0pUHAE1jRBeSRhMW7ogXvDHEdeaZGioNE4,702
-brainstate/event/_csr_mv.py,sha256=HStHvK3KyEMfLsIUslZjgbdU6OsD1yKGrzQOzBXG36M,10266
-brainstate/event/_csr_test.py,sha256=_iXwUFq90GU7npVOUnlI4NA27RJ8zyCZBxe7NDH803o,9533
-brainstate/event/_fixedprob_mv.py,sha256=nR3lhd87t1Vge435QHnFuDp-UBbWoW0Qk1kbsjRHQyc,25541
-brainstate/event/_fixedprob_mv_benchmark.py,sha256=_F_8fH5MNMJZHeSqnq9DYMI9OgYr6JIxBKjbsgeWRv4,4720
-brainstate/event/_fixedprob_mv_test.py,sha256=pVEarvGbqTjnAbxgMVRTAhkyYbvDnlyCJdeOdDD927w,4283
-brainstate/event/_linear_mv.py,sha256=O5qbY31GNV1qEDrZ5kvPbA8Ae-bY5JpUgGtqDFNAeV0,11794
-brainstate/event/_linear_mv_benckmark.py,sha256=hu0WqYMIa3jMoH7Fq9dgxcBjjXGFhghPx9vztyCo1KY,2411
-brainstate/event/_linear_mv_test.py,sha256=V9w41ZP2vu95CyCdCkm-j9Eftqs2kqmeBY809N1-syY,3736
-brainstate/event/_misc.py,sha256=8IpPooXjF2m0-tuo3pGHqThq2yLSNmYziy_zdurZ3NI,1040
-brainstate/event/_xla_custom_op.py,sha256=wF_nKgLUv1IGd8OY89MYqIvyZITl8UcrVysJWFugJxY,11093
-brainstate/event/_xla_custom_op_test.py,sha256=rnkGMleXzLfJj4y5QqwfBvCCLTAHe_uabwBDniY-URM,1745
 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
 brainstate/functional/_normalization.py,sha256=i2EV7hSsqcNdcYRX2wAxjq8doHwyN9eNJTGTaPt03xE,2605
 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=fyvQMlAUY3QYTzvDzz5TDoWS2XQwZ6P3ic6BtysZyHM,1026
+brainstate/graph/__init__.py,sha256=noo4TjBg6iEhjjwk0sAGUhR7Ge-z8Vnc2rLYUvnqttw,1295
 brainstate/graph/_graph_node.py,sha256=swAokZLKswSTaq2WEhyLIs38sy_67C6maHI6T3e1hvY,8339
 brainstate/graph/_graph_node_test.py,sha256=BFGfdzZFDHI0XK7hHotSVWKt3em1taGvn8FHF9NCXx8,2702
 brainstate/graph/_graph_operation.py,sha256=cIwGo3ICgtce2fmdn917r81evMFjJIKeW9doaQK4DD8,64111
@@ -79,7 +65,7 @@ brainstate/nn/_dyn_impl/_dynamics_neuron.py,sha256=cTbIn41EPYG0h3ICzKBXxpgB6wwA2
 brainstate/nn/_dyn_impl/_dynamics_neuron_test.py,sha256=Tfzrzu7udGrLJGnqItiLWe5WT0dgduvYOgzGCnaPJQg,6317
 brainstate/nn/_dyn_impl/_dynamics_synapse.py,sha256=MsbPyaiDyjungyzuK2b3exRGaMpZgmsmmNHNLjgxQKw,15269
 brainstate/nn/_dyn_impl/_dynamics_synapse_test.py,sha256=t5i-HV0ii9sUNzWTEv04o26QVtQ-mCdMJcFq2MD755A,4981
-brainstate/nn/_dyn_impl/_inputs.py,sha256=UNoGxKIKXwPnhelljDowqAWlV6ds7aBBkEbvdy2oDI4,11302
+brainstate/nn/_dyn_impl/_inputs.py,sha256=QOUpAb2YJOE78uAvIS8Ep6MFcQHV-V6uRwmYvk5p9bk,11385
 brainstate/nn/_dyn_impl/_projection_alignpost.py,sha256=PNC1Tzx_SF2DHAHeJCufXzO_Q4qLoBpWABI45B3GRuc,876
 brainstate/nn/_dyn_impl/_rate_rnns.py,sha256=dz_yT_6hJVhKulcjIARbGtmMzZqISws96CtBc6o5GOo,14768
 brainstate/nn/_dyn_impl/_rate_rnns_test.py,sha256=gNgtr-a4ZiU1XF9wFG1HiJ9fLosfWchVR9Zn1x39xt4,2452
@@ -94,7 +80,7 @@ brainstate/nn/_dynamics/_synouts.py,sha256=9TGAc-nVa50th7KKn4oKLbro-4W4rwxYvp-eu
 brainstate/nn/_dynamics/_synouts_test.py,sha256=V_jDswRN4VvEXD-2yJO3VA1TALgX0HK6oPBQiUntOWc,2266
 brainstate/nn/_elementwise/__init__.py,sha256=PK8oq1K_EG2941AiUyLxCWoRdWvMO3yt8ZJbw3Lkhu8,935
 brainstate/nn/_elementwise/_dropout.py,sha256=0Ebo-2y1VswvBqZ7sCA0SEUm37y49EUsef8oiSFpYGk,17759
-brainstate/nn/_elementwise/_dropout_test.py,sha256=ZzNvjFf46NpKWGBIcT6O0lKOBGpxOStOAIGM4cE8LfE,4405
+brainstate/nn/_elementwise/_dropout_test.py,sha256=k6aB5v8RYMoV5w8UV9UNSFhaQTV7woS6jx3SNESuCRs,4383
 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
@@ -131,8 +117,8 @@ brainstate/util/_others.py,sha256=jsPZwP-v_5HRV-LB5F0NUsiqr04y8bmGIsu_JMyVcbQ,14
 brainstate/util/_pretty_repr.py,sha256=bDpU4gbkS4B8cXBkiN8kBQNmruxiJzDRF-eIqzyeYnM,5716
 brainstate/util/_scaling.py,sha256=pc_eM_SZVwkY65I4tJh1ODiHNCoEhsfFXl2zBK0PLAg,7562
 brainstate/util/_struct.py,sha256=KMMHcshOM20gYhSahNzWLxsTt-Rt3AeX3Uz26-rP9vI,17619
-brainstate-0.1.0.post20250120.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.1.0.post20250120.dist-info/METADATA,sha256=vUyr4XjiyAW68waFKMray9EEFHTqjqRp5GlqAG8LsKY,3585
-brainstate-0.1.0.post20250120.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.1.0.post20250120.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.1.0.post20250120.dist-info/RECORD,,
+brainstate-0.1.0.post20250127.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.1.0.post20250127.dist-info/METADATA,sha256=-j8gDuJN37nbSkijFC9TtG5UbflMZZWK-sftY7fkeps,3585
+brainstate-0.1.0.post20250127.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+brainstate-0.1.0.post20250127.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.1.0.post20250127.dist-info/RECORD,,

brainstate/event/__init__.py

@@ -1,27 +0,0 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-
-from ._csr import *
-from ._fixedprob_mv import *
-from ._linear_mv import *
-from ._xla_custom_op import *
-
-__all__ = [
-    'FixedProb',
-    'XLACustomOp',
-    'CSR',
-    'CSC',
-]