brainstate 0.0.1.post20240612__py2.py3-none-any.whl → 0.0.1.post20240622__py2.py3-none-any.whl

brainstate/functional/_spikes.py

@@ -87,4 +87,3 @@ def spike_bitwise(x, y, op: str):
     return spike_bitwise_ixor(x, y)
   else:
     raise NotImplementedError(f"Unsupported bitwise operation: {op}.")
-

brainstate/mixin.py

@@ -68,7 +68,7 @@ class DelayedInit(Mixin):
   Note this Mixin can be applied in any Python object.
   """
 
-  non_hash_params: Optional[Sequence[str]] = None
+  non_hashable_params: Optional[Sequence[str]] = None
 
   @classmethod
   def delayed(cls, *args, **kwargs) -> 'DelayedInitializer':
@@ -94,7 +94,7 @@ class DelayedInitializer(metaclass=NoSubclassMeta):
   """
 
   def __init__(self, cls: T, *desc_tuple, **desc_dict):
-    self.cls = cls
+    self.cls: type = cls
 
     # arguments
     self.args = desc_tuple

brainstate/nn/_elementwise.py

@@ -19,11 +19,12 @@ from __future__ import annotations
 
 from typing import Optional
 
+import brainunit as bu
 import jax.numpy as jnp
 import jax.typing
 
 from ._base import ElementWiseBlock
-from .. import math, environ, random, functional as F
+from .. import environ, random, functional as F
 from .._module import Module
 from .._state import ParamState
 from ..mixin import Mode
@@ -82,7 +83,7 @@ class Threshold(Module, ElementWiseBlock):
     self.value = value
 
   def __call__(self, x: ArrayLike) -> ArrayLike:
-    dtype = math.get_dtype(x)
+    dtype = bu.math.get_dtype(x)
     return jnp.where(x > jnp.asarray(self.threshold, dtype=dtype),
                      x,
                      jnp.asarray(self.value, dtype=dtype))
@@ -1142,7 +1143,7 @@ class Dropout(Module, ElementWiseBlock):
     self.prob = prob
 
   def __call__(self, x):
-    dtype = math.get_dtype(x)
+    dtype = bu.math.get_dtype(x)
     fit_phase = environ.get('fit', desc='Whether this is a fitting process. Bool.')
     if fit_phase:
       keep_mask = random.bernoulli(self.prob, x.shape)
@@ -1172,7 +1173,7 @@ class _DropoutNd(Module, ElementWiseBlock):
     self.channel_axis = channel_axis
 
   def __call__(self, x):
-    dtype = math.get_dtype(x)
+    dtype = bu.math.get_dtype(x)
     # get fit phase
     fit_phase = environ.get('fit', desc='Whether this is a fitting process. Bool.')
 

brainstate/nn/_misc.py

@@ -20,9 +20,10 @@ from enum import Enum
 from functools import wraps
 from typing import Sequence, Callable
 
+import brainunit as bu
 import jax.numpy as jnp
 
-from .. import environ, math
+from .. import environ
 from .._state import State
 from ..transform import vector_grad
 
@@ -96,7 +97,7 @@ def exp_euler(fun):
       )
     dt = environ.get('dt')
     linear, derivative = vector_grad(fun, argnums=0, return_value=True)(*args, **kwargs)
-    phi = math.exprel(dt * linear)
+    phi = bu.math.exprel(dt * linear)
     return args[0] + dt * phi * derivative
 
   return integral
@@ -128,5 +129,5 @@ def exp_euler_step(fun: Callable, *args, **kwargs):
     )
   dt = environ.get('dt')
   linear, derivative = vector_grad(fun, argnums=0, return_value=True)(*args, **kwargs)
-  phi = math.exprel(dt * linear)
+  phi = bu.math.exprel(dt * linear)
   return args[0] + dt * phi * derivative

brainstate/nn/_others.py

@@ -19,10 +19,11 @@ from __future__ import annotations
 from functools import partial
 from typing import Optional
 
+import brainunit as bu
 import jax.numpy as jnp
 
 from ._base import DnnLayer
-from .. import random, math, environ, typing, init
+from .. import random, environ, typing, init
 from ..mixin import Mode
 
 __all__ = [
@@ -88,7 +89,7 @@ class DropoutFixed(DnnLayer):
     self.mask = init.param(partial(random.bernoulli, self.prob), self.in_size, batch_size)
 
   def update(self, x):
-    dtype = math.get_dtype(x)
+    dtype = bu.math.get_dtype(x)
     fit_phase = environ.get('fit', desc='Whether this is a fitting process. Bool.')
     if fit_phase:
       assert self.mask.shape == x.shape, (f"Input shape {x.shape} does not match the mask shape {self.mask.shape}. "

brainstate/nn/_poolings.py

@@ -21,12 +21,13 @@ import functools
 from typing import Sequence, Optional
 from typing import Union, Tuple, Callable, List
 
+import brainunit as bu
 import jax
 import jax.numpy as jnp
 import numpy as np
 
 from ._base import DnnLayer, ExplicitInOutSize
-from .. import environ, math
+from .. import environ
 from ..mixin import Mode
 from ..typing import Size
 
@@ -53,8 +54,8 @@ class Flatten(DnnLayer, ExplicitInOutSize):
 
   Args:
       in_size: Sequence of int. The shape of the input tensor.
-      start_dim: first dim to flatten (default = 1).
-      end_dim: last dim to flatten (default = -1).
+      start_axis: first dim to flatten (default = 1).
+      end_axis: last dim to flatten (default = -1).
 
   Examples::
       >>> import brainstate as bst
@@ -74,36 +75,36 @@ class Flatten(DnnLayer, ExplicitInOutSize):
 
   def __init__(
       self,
-      start_dim: int = 0,
-      end_dim: int = -1,
+      start_axis: int = 0,
+      end_axis: int = -1,
       in_size: Optional[Size] = None
   ) -> None:
     super().__init__()
-    self.start_dim = start_dim
-    self.end_dim = end_dim
+    self.start_axis = start_axis
+    self.end_axis = end_axis
 
     if in_size is not None:
       self.in_size = tuple(in_size)
-      y = jax.eval_shape(functools.partial(math.flatten, start_dim=start_dim, end_dim=end_dim),
+      y = jax.eval_shape(functools.partial(bu.math.flatten, start_axis=start_axis, end_axis=end_axis),
                          jax.ShapeDtypeStruct(self.in_size, environ.dftype()))
       self.out_size = y.shape
 
   def update(self, x):
     if self._in_size is None:
-      start_dim = self.start_dim if self.start_dim >= 0 else x.ndim + self.start_dim
+      start_axis = self.start_axis if self.start_axis >= 0 else x.ndim + self.start_axis
     else:
       assert x.ndim >= len(self.in_size), 'Input tensor has fewer dimensions than the expected shape.'
       dim_diff = x.ndim - len(self.in_size)
       if self.in_size != x.shape[dim_diff:]:
         raise ValueError(f'Input tensor has shape {x.shape}, but expected shape {self.in_size}.')
-      if self.start_dim >= 0:
-        start_dim = self.start_dim + dim_diff
+      if self.start_axis >= 0:
+        start_axis = self.start_axis + dim_diff
       else:
-        start_dim = x.ndim + self.start_dim
-    return math.flatten(x, start_dim, self.end_dim)
+        start_axis = x.ndim + self.start_axis
+    return bu.math.flatten(x, start_axis, self.end_axis)
 
   def __repr__(self) -> str:
-    return f'{self.__class__.__name__}(start_dim={self.start_dim}, end_dim={self.end_dim})'
+    return f'{self.__class__.__name__}(start_axis={self.start_axis}, end_axis={self.end_axis})'
 
 
 class Unflatten(DnnLayer, ExplicitInOutSize):
@@ -124,7 +125,7 @@ class Unflatten(DnnLayer, ExplicitInOutSize):
         :math:`\prod_{i=1}^n U_i = S_{\text{dim}}`.
 
   Args:
-      dim: int, Dimension to be unflattened.
+      axis: int, Dimension to be unflattened.
       sizes: Sequence of int. New shape of the unflattened dimension.
       in_size: Sequence of int. The shape of the input tensor.
   """
@@ -132,7 +133,7 @@ class Unflatten(DnnLayer, ExplicitInOutSize):
 
   def __init__(
       self,
-      dim: int,
+      axis: int,
       sizes: Size,
       mode: Mode = None,
       name: str = None,
@@ -140,7 +141,7 @@ class Unflatten(DnnLayer, ExplicitInOutSize):
   ) -> None:
     super().__init__(mode=mode, name=name)
 
-    self.dim = dim
+    self.axis = axis
     self.sizes = sizes
     if isinstance(sizes, (tuple, list)):
       for idx, elem in enumerate(sizes):
@@ -152,15 +153,15 @@ class Unflatten(DnnLayer, ExplicitInOutSize):
 
     if in_size is not None:
       self.in_size = tuple(in_size)
-      y = jax.eval_shape(functools.partial(math.unflatten, dim=dim, sizes=sizes),
+      y = jax.eval_shape(functools.partial(bu.math.unflatten, axis=axis, sizes=sizes),
                          jax.ShapeDtypeStruct(self.in_size, environ.dftype()))
       self.out_size = y.shape
 
   def update(self, x):
-    return math.unflatten(x, self.dim, self.sizes)
+    return bu.math.unflatten(x, self.axis, self.sizes)
 
   def __repr__(self):
-    return f'{self.__class__.__name__}(dim={self.dim}, sizes={self.sizes})'
+    return f'{self.__class__.__name__}(axis={self.axis}, sizes={self.sizes})'
 
 
 class _MaxPool(DnnLayer, ExplicitInOutSize):

brainstate/nn/_poolings_test.py

@@ -18,7 +18,7 @@ class TestFlatten(parameterized.TestCase):
       (10, 20, 30),
     ]:
       arr = bst.random.rand(*size)
-      f = nn.Flatten(start_dim=0)
+      f = nn.Flatten(start_axis=0)
       out = f(arr)
       self.assertTrue(out.shape == (np.prod(size),))
 
@@ -29,21 +29,21 @@ class TestFlatten(parameterized.TestCase):
       (10, 20, 30),
     ]:
       arr = bst.random.rand(*size)
-      f = nn.Flatten(start_dim=1)
+      f = nn.Flatten(start_axis=1)
       out = f(arr)
       self.assertTrue(out.shape == (size[0], np.prod(size[1:])))
 
   def test_flatten3(self):
     size = (16, 32, 32, 8)
     arr = bst.random.rand(*size)
-    f = nn.Flatten(start_dim=0, in_size=(32, 8))
+    f = nn.Flatten(start_axis=0, in_size=(32, 8))
     out = f(arr)
     self.assertTrue(out.shape == (16, 32, 32 * 8))
 
   def test_flatten4(self):
     size = (16, 32, 32, 8)
     arr = bst.random.rand(*size)
-    f = nn.Flatten(start_dim=1, in_size=(32, 32, 8))
+    f = nn.Flatten(start_axis=1, in_size=(32, 32, 8))
     out = f(arr)
     self.assertTrue(out.shape == (16, 32, 32 * 8))
 

brainstate/optim/__init__.py

@@ -20,4 +20,3 @@ from ._sgd_optimizer import *
 from ._sgd_optimizer import __all__ as optimizer_all
 
 __all__ = scheduler_all + optimizer_all
-

brainstate/optim/_sgd_optimizer.py

@@ -18,11 +18,12 @@
 import functools
 from typing import Union, Dict, Optional, Tuple, Any, TypeVar
 
+import brainunit as bu
 import jax
 import jax.numpy as jnp
 
 from ._lr_scheduler import make_schedule, LearningRateScheduler
-from .. import environ, math
+from .. import environ
 from .._module import Module
 from .._state import State, LongTermState, StateDictManager, visible_state_dict
 
@@ -282,7 +283,7 @@ class Momentum(_WeightDecayOptimizer):
     for k, v in train_states.items():
       assert isinstance(v, State), f'"{k}" must be an instance of brainstate.State.'
       self.weight_states.add_unique_elem(k, v)
-      self.momentum_states[k] = OptimState(math.tree_zeros_like(v.value))
+      self.momentum_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
 
   def update(self, grads: dict):
     lr = self.lr()
@@ -349,7 +350,7 @@ class MomentumNesterov(_WeightDecayOptimizer):
     for k, v in train_states.items():
       assert isinstance(v, State), f'"{k}" must be an instance of brainstate.State.'
       self.weight_states.add_unique_elem(k, v)
-      self.momentum_states[k] = OptimState(math.tree_zeros_like(v.value))
+      self.momentum_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
 
   def update(self, grads: dict):
     lr = self.lr()
@@ -417,7 +418,7 @@ class Adagrad(_WeightDecayOptimizer):
     for k, v in train_states.items():
       assert isinstance(v, State), f'"{k}" must be an instance of brainstate.State.'
       self.weight_states.add_unique_elem(k, v)
-      self.cache_states[k] = OptimState(math.tree_zeros_like(v.value))
+      self.cache_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
 
   def update(self, grads: dict):
     lr = self.lr()
@@ -500,8 +501,8 @@ class Adadelta(_WeightDecayOptimizer):
     for k, v in train_states.items():
       assert isinstance(v, State), f'"{k}" must be an instance of brainstate.State.'
       self.weight_states.add_unique_elem(k, v)
-      self.cache_states[k] = OptimState(math.tree_zeros_like(v.value))
-      self.delta_states[k] = OptimState(math.tree_zeros_like(v.value))
+      self.cache_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
+      self.delta_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
 
   def update(self, grads: dict):
     weight_values, grad_values, cache_values, delta_values = to_same_dict_tree(
@@ -574,7 +575,7 @@ class RMSProp(_WeightDecayOptimizer):
     for k, v in train_states.items():
       assert isinstance(v, State), f'"{k}" must be an instance of brainstate.State.'
       self.weight_states.add_unique_elem(k, v)
-      self.cache_states[k] = OptimState(math.tree_zeros_like(v.value))
+      self.cache_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
 
   def update(self, grads: dict):
     lr = self.lr()
@@ -647,8 +648,8 @@ class Adam(_WeightDecayOptimizer):
     for k, v in train_states.items():
       assert isinstance(v, State), f'"{k}" must be an instance of brainstate.State.'
       self.weight_states.add_unique_elem(k, v)
-      self.m1_states[k] = OptimState(math.tree_zeros_like(v.value))
-      self.m2_states[k] = OptimState(math.tree_zeros_like(v.value))
+      self.m1_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
+      self.m2_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
 
   def update(self, grads: dict):
     lr = self.lr()
@@ -730,7 +731,7 @@ class LARS(_WeightDecayOptimizer):
     for k, v in train_states.items():
       assert isinstance(v, State), f'"{k}" must be an instance of brainstate.State.'
       self.weight_states.add_unique_elem(k, v)
-      self.momentum_states[k] = OptimState(math.tree_zeros_like(v.value))
+      self.momentum_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
 
   def update(self, grads: dict):
     lr = self.lr()
@@ -835,10 +836,10 @@ class Adan(_WeightDecayOptimizer):
     for k, v in train_states.items():
       assert isinstance(v, State), f'"{k}" must be an instance of brainstate.State.'
       self.weight_states.add_unique_elem(k, v)
-      self.exp_avg_states[k] = OptimState(math.tree_zeros_like(v.value))
-      self.exp_avg_sq_states[k] = OptimState(math.tree_zeros_like(v.value))
-      self.exp_avg_diff_states[k] = OptimState(math.tree_zeros_like(v.value))
-      self.pre_grad_states[k] = OptimState(math.tree_zeros_like(v.value))
+      self.exp_avg_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
+      self.exp_avg_sq_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
+      self.exp_avg_diff_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
+      self.pre_grad_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
 
   def update(self, grads: dict):
     lr = self.lr()
@@ -989,10 +990,10 @@ class AdamW(_WeightDecayOptimizer):
     for k, v in train_states.items():
       assert isinstance(v, State), f'"{k}" must be an instance of brainstate.State.'
       self.weight_states.add_unique_elem(k, v)
-      self.m1_states[k] = OptimState(math.tree_zeros_like(v.value))
-      self.m2_states[k] = OptimState(math.tree_zeros_like(v.value))
+      self.m1_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
+      self.m2_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
       if self.amsgrad:
-        self.vmax_states[k] = OptimState(math.tree_zeros_like(v.value))
+        self.vmax_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
 
   def update(self, grads: dict):
     lr_old = self.lr()

brainstate/transform/__init__.py

@@ -17,10 +17,10 @@
 This module contains the functions for the transformation of the brain data.
 """
 
-from ._control import *
-from ._control import __all__ as _controls_all
 from ._autograd import *
 from ._autograd import __all__ as _gradients_all
+from ._control import *
+from ._control import __all__ as _controls_all
 from ._jit import *
 from ._jit import __all__ as _jit_all
 from ._jit_error import *
@@ -33,4 +33,3 @@ from ._progress_bar import __all__ as _progress_bar_all
 __all__ = _gradients_all + _jit_error_all + _controls_all + _make_jaxpr_all + _jit_all + _progress_bar_all
 
 del _gradients_all, _jit_error_all, _controls_all, _make_jaxpr_all, _jit_all, _progress_bar_all
-

brainstate/transform/_autograd.py

@@ -25,8 +25,8 @@ from jax._src.api import _vjp
 from jax.api_util import argnums_partial
 from jax.extend import linear_util
 
-from brainstate._utils import set_module_as
 from brainstate._state import State, StateTrace, StateDictManager
+from brainstate._utils import set_module_as
 
 __all__ = [
   'vector_grad', 'grad', 'jacrev', 'jacfwd', 'jacobian', 'hessian',

brainstate/transform/_autograd_test.py

@@ -537,7 +537,6 @@ class TestPureFuncJacobian(unittest.TestCase):
 
   def test_jacrev_return_aux1(self):
     with bc.environ.context(precision=64):
-
       def f1(x, y):
         a = 4 * x[1] ** 2 - 2 * x[2]
         r = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1], a, x[2] * jnp.sin(x[0])])
@@ -564,7 +563,6 @@ class TestPureFuncJacobian(unittest.TestCase):
       assert (vec == _r).all()
 
 
-
 class TestClassFuncJacobian(unittest.TestCase):
   def test_jacrev1(self):
     def f1(x, y):

brainstate/transform/_jit_test.py

@@ -16,7 +16,6 @@
 import unittest
 
 import jax.numpy as jnp
-import jax.stages
 
 import brainstate as bc
 
@@ -90,7 +89,6 @@ class TestJIT(unittest.TestCase):
     self.assertTrue(len(compiling) == 2)
 
   def test_jit_attribute_origin_fun(self):
-
     def fun1(x):
       return x
 
@@ -99,4 +97,3 @@ class TestJIT(unittest.TestCase):
     self.assertTrue(isinstance(jitted_fun.stateful_fun, bc.transform.StatefulFunction))
     self.assertTrue(callable(jitted_fun.jitted_fun))
     self.assertTrue(callable(jitted_fun.clear_cache))
-

brainstate/transform/_make_jaxpr.py

@@ -75,7 +75,6 @@ from brainstate._utils import set_module_as
 PyTree = Any
 AxisName = Hashable
 
-
 __all__ = [
   "StatefulFunction",
   "make_jaxpr",

brainstate/transform/_make_jaxpr_test.py

@@ -129,5 +129,3 @@ def test_return_states():
 
   with pytest.raises(ValueError):
     f()
-
-

brainstate/transform/_progress_bar.py

@@ -14,13 +14,12 @@
 # ==============================================================================
 
 from __future__ import annotations
+
 import copy
 from typing import Optional
 
 import jax
 
-from brainstate import environ
-
 try:
   from tqdm.auto import tqdm
 except (ImportError, ModuleNotFoundError):
@@ -95,7 +94,6 @@ class ProgressBarRunner(object):
     self.tqdm_bars[0].close()
 
   def __call__(self, iter_num, *args, **kwargs):
-
     _ = jax.lax.cond(
       iter_num == 0,
       lambda: jax.debug.callback(self._define_tqdm),

brainstate/util.py

@@ -26,7 +26,6 @@ from jax.lib import xla_bridge
 
 from ._utils import set_module_as
 
-
 __all__ = [
   'unique_name',
   'clear_buffer_memory',

{brainstate-0.0.1.post20240612.dist-info → brainstate-0.0.1.post20240622.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.0.1.post20240612
+Version: 0.0.1.post20240622
 Summary: A State-based Transformation System for Brain Dynamics Programming.
 Home-page: https://github.com/brainpy/brainstate
 Author: BrainPy Team
@@ -31,21 +31,11 @@ License-File: LICENSE
 Requires-Dist: jax
 Requires-Dist: jaxlib
 Requires-Dist: numpy
+Requires-Dist: brainunit
 Provides-Extra: cpu
 Requires-Dist: jaxlib ; extra == 'cpu'
-Requires-Dist: brainpylib ; extra == 'cpu'
-Provides-Extra: cpu_mini
-Requires-Dist: jaxlib ; extra == 'cpu_mini'
-Provides-Extra: cuda11
-Requires-Dist: jaxlib[cuda11_pip] ; extra == 'cuda11'
-Requires-Dist: brainpylib ; extra == 'cuda11'
-Provides-Extra: cuda11_mini
-Requires-Dist: jaxlib[cuda11_pip] ; extra == 'cuda11_mini'
 Provides-Extra: cuda12
 Requires-Dist: jaxlib[cuda12_pip] ; extra == 'cuda12'
-Requires-Dist: brainpylib ; extra == 'cuda12'
-Provides-Extra: cuda12_mini
-Requires-Dist: jaxlib[cuda12_pip] ; extra == 'cuda12_mini'
 Provides-Extra: testing
 Requires-Dist: pytest ; extra == 'testing'
 Provides-Extra: tpu

{brainstate-0.0.1.post20240612.dist-info → brainstate-0.0.1.post20240622.dist-info}/RECORD

@@ -1,45 +1,38 @@
-brainstate/__init__.py,sha256=3R0I9oLpIS6Z9iwmx5ODzSlZGX3MbYWvCWsFHaCiaG4,1436
-brainstate/_module.py,sha256=R3pBeNvqR_mEfquZU60uWj7JmopOCUciF2BgcJyA0aw,48151
-brainstate/_module_test.py,sha256=4tqtp2-j5mSoUmCITY0mVZEcXzxXCWJ_02Jdt1fxYJg,4502
+brainstate/__init__.py,sha256=DwgnJOghZ_qeFh0a_roiaMCDH-V_F6Ve7by3xjSVrwk,1408
+brainstate/_module.py,sha256=lIfRRev49QQfq55PBm9YohtvPsMWdLHNhrfBjMjD6c8,52623
+brainstate/_module_test.py,sha256=TJlxR4R5bf621y68hTgzTaf0PBN9YmVhwoGKNcpXbpE,7821
 brainstate/_state.py,sha256=RWnLjMeaidxWXNAA0X-8mxj4i61j3T8w5KhugACUYhI,11422
 brainstate/_state_test.py,sha256=HDdipndRLhEHWEdTmyT1ayEBkbv6qJKykfCWKI6yJ_E,1253
 brainstate/_utils.py,sha256=RLorgGJkt2BhbX4C-ygd-PPG0wfcGCghjSP93sRvzqM,833
-brainstate/environ.py,sha256=RMDUACuixwk2ZTHf0UGLhcd5DCraW-l1j9T3wc2wcFc,10242
-brainstate/mixin.py,sha256=V75vjMTzYcCMlPo5wekgRZZ9o6-xN8kJocQgEliu5yI,10738
+brainstate/environ.py,sha256=LwRwnFaTbv8l7nHRIbSV46WzcN7pGLQFhT_xDUox2yA,10240
+brainstate/mixin.py,sha256=x4WIYMTCFZgtTp-uiZeNI5J4Qd2BYaV0Ccm_EMdzl9c,10748
 brainstate/mixin_test.py,sha256=qDYqhHbHw3aBFW8aHQdPhov29013Eo9TJDF7RW2dapE,2919
 brainstate/random.py,sha256=Mi5i0kAsR8C-VoI8LMuIbPPr6YFzq6NBxhJ5K0w2qW4,186392
 brainstate/random_test.py,sha256=cCeuYvlZkCS2_RgG0vipZFNSHG8b-uJ7SXM9SZDCYQM,17866
 brainstate/surrogate.py,sha256=1kgbn82GSlpReIytIVl29yozk75gkdZv0gTBlixQ4C4,43798
 brainstate/typing.py,sha256=Ooweu7c17nYP686fyIeKNomChodSxx_OEpu8QRoB9cY,2180
-brainstate/util.py,sha256=FrBN_OZAPlWxfNK8c9Z1d-bbIa8qwMrcOsSJZJS8xOE,19878
-brainstate/functional/__init__.py,sha256=Z-43coOHFAsQK0u5amlr4l0fNNPc7dVcuKXfNY4Gj_s,1107
-brainstate/functional/_activations.py,sha256=IfZ6Zy8SAwyxo166E3NmCZMUHnG_rBFAUaLTyxG5FgA,18490
+brainstate/util.py,sha256=y-6eX1z3EMyg6pfZt4YdDalOnJ3HDAT1IPBCJDp-gQI,19876
+brainstate/functional/__init__.py,sha256=j6-3Er4fgqWpvntzYCZVB3e5hoz-Z3aqvapITCuDri0,1107
+brainstate/functional/_activations.py,sha256=xlwvYG8qvpkfMEZTFxD_4amW63ZfEa8x3vzVH2hDgeY,17791
 brainstate/functional/_normalization.py,sha256=IxE580waloZylZVXcpUUK4bWQdlE6oSPfafaKYfDkbg,2169
 brainstate/functional/_others.py,sha256=ifB-l82y7ZB632yLUJOEcpkRY-yOoiJ0mtDOxNilp4M,1711
-brainstate/functional/_spikes.py,sha256=uAln_Q87pr1codLxeDck3PUA9jpk7S5LifNps1kdyrU,2576
+brainstate/functional/_spikes.py,sha256=70qGvo4B--QtxfJMjLwGmk9pVsf2x2YNEEgjT-il_Jw,2574
 brainstate/init/__init__.py,sha256=R1dHgub47o-WJM9QkFLc7x_Q7GsyaKKDtrRHTFPpC5g,1097
 brainstate/init/_base.py,sha256=jRTmfoUsH_315vW9YMZzyIn2KDAAsv56SplBnvOyBW0,1148
 brainstate/init/_generic.py,sha256=OJFS7DHYmZV0JogdsgjnUseUfvTUrAUYiXZynCQqmG4,5163
 brainstate/init/_random_inits.py,sha256=STbX-mrHwNuICXkw7EldtJLdUUsWOAcGkEzx2ycV-Yc,15321
 brainstate/init/_regular_inits.py,sha256=n-vF-51FM1UcUh-8h5lUk5Jhjrn04KPcGXgGhUGFAAk,3065
-brainstate/math/__init__.py,sha256=meQnO6k1EzMRMhO3x_22oj4-LVo_KevHK4L04bmHZPo,873
-brainstate/math/_einops.py,sha256=Lwi8AGKNPb-x1To0dDQYHbKwUOrO6pPL23qdg28-nB0,31726
-brainstate/math/_einops_parsing.py,sha256=zjTJdJlEBRS0y02PgKoZ8Y6bv54B4Axzk4AtPQOo934,6805
-brainstate/math/_einops_parsing_test.py,sha256=JPn73yld300481J6E9cL7jHWn63Vr21VV8k1jJxAK4A,4888
-brainstate/math/_einops_test.py,sha256=xj-DDTL0EsW1Obm64KCnT7eqELWjjj04Ozdwk0839Tw,13289
-brainstate/math/_misc.py,sha256=jDtREP4ojxHyj6lXcLcYLGVsLA0HFZcrs8cdlnA7aK8,7863
-brainstate/math/_misc_test.py,sha256=V41YV-RiEbukKQlzq54174cpSalOhMjaHOoVH8o82eI,2443
 brainstate/nn/__init__.py,sha256=YJHoI8cXKVRS8f2vUl3Zegp5wm0svMz3qo9JmQJiMQk,2162
 brainstate/nn/_base.py,sha256=lzbZpku3Q2arH6ZaAwjs6bhbV0RcFChxo2UcpnX5t84,8481
 brainstate/nn/_connections.py,sha256=GSOW2IbpJRHdPyF4nFJ2RPgO8y6SVHT1Gn-pbri9pMk,22970
 brainstate/nn/_dynamics.py,sha256=OeYYXv1dqjUDcCsRhZo1XS7SP2li1vlH9uhME_PE9v0,13205
-brainstate/nn/_elementwise.py,sha256=T1oCu47t11Ki7LPaL-hHk4W8bKP_Q3HLJcGngcmGK0Y,43552
+brainstate/nn/_elementwise.py,sha256=6BTqSvSnaHhldwB5ol5OV0hPJ5yJ-Jpm4WSrtFKMNoQ,43579
 brainstate/nn/_embedding.py,sha256=WbgrIaM_14abN8zBDr0xipBOsFc8dXP2m7Z_aRLAfmU,2249
-brainstate/nn/_misc.py,sha256=Z7gdJraJ18gVNHyNOk_KmE67M3OM4z3QT4RN6al5JMc,3766
+brainstate/nn/_misc.py,sha256=Xc4U4NLmvfnKdBNDayFrRBPAy3p0beS6T9C59rIDP00,3790
 brainstate/nn/_normalizations.py,sha256=9yVDORAEpqEkL9MYSPU4m7C4q8Qj5UNsPh9sKmIt5gQ,14329
-brainstate/nn/_others.py,sha256=8PYmCiUNzru4kmm58HY0RzCs-32dnwNFDZdTTPixaqo,4492
-brainstate/nn/_poolings.py,sha256=cNZ1PyMIaViP-_AUkEbpy3ZfHo--ib1hAhL0bEAmXIQ,45688
-brainstate/nn/_poolings_test.py,sha256=iE0NgvOIWVgwmcvP4wazhGG4RJQdU2eeagdJ1sDXIBQ,7260
+brainstate/nn/_others.py,sha256=AYyrbbdKZj16kT0cVITnoZHck4xcccM1W3LX5XM5Z3Q,4513
+brainstate/nn/_poolings.py,sha256=wO1Q4s8blsLLv4CMlkrvZm0ravdL3dFGyOcg2QDendI,45754
+brainstate/nn/_poolings_test.py,sha256=Mj4gO86Xl4JS5hHNR_CgeUdZQIqAxUoeBldS-eoZoBg,7264
 brainstate/nn/_rate_rnns.py,sha256=Cebhy57UWzfwrCfq0v2qLDegmb__mXL5ht750y4aTro,14457
 brainstate/nn/_readout.py,sha256=jsQwhVnrJICKw4wFq-Du2AORPb_XXz_tZ4cURcckU-E,4240
 brainstate/nn/_synouts.py,sha256=gi3EyKlzt4UoyghwvNIr03r7YabZyl1idbq9aYG8zYM,4379
@@ -49,23 +42,23 @@ brainstate/nn/_projection/_align_pre.py,sha256=R2U6_RQ_o8y6PWXpozeWE2cx_oQ7WMhhr
 brainstate/nn/_projection/_delta.py,sha256=KT8ySo3n_Q_7swzOH-ISDf0x9rjMkiv99H-vqeQZDR8,7122
 brainstate/nn/_projection/_utils.py,sha256=UcmELOqsINgqJr7eC5BSNNteyZ--1lyGjhUTJfxyMmA,813
 brainstate/nn/_projection/_vanilla.py,sha256=_bh_DLtF0o33SBtj6IGL8CTanFEtJwfjBrgxBEAmIlg,3397
-brainstate/optim/__init__.py,sha256=1xH5_peSWKuZ4tOU295r9EKAv0a-cBMABx6XV3faDJI,919
+brainstate/optim/__init__.py,sha256=1L6x_qZprw3PJYddB1nX-uTFGUl6_Qt3PM0OdY6g968,917
 brainstate/optim/_lr_scheduler.py,sha256=emKnA52UVqOfUcX7LJqwP-FVDVlGGzTQi2djYmbCWUo,15627
 brainstate/optim/_lr_scheduler_test.py,sha256=OwF8Iz-PorEbO0gO--A7IIgQEytqEfYWbPucAgzqL90,1598
-brainstate/optim/_sgd_optimizer.py,sha256=7-jMfP_Hol0XGEA6_4wVqygpLTqI1646F6eeLtwtNFY,45760
-brainstate/transform/__init__.py,sha256=9S9TLp1sF6nWRmW6jFtu6_dLmOc43V88Ruh073Z8I50,1460
-brainstate/transform/_autograd.py,sha256=sFGJ6oAhlSr54Hb1c1aNc5Q2St7eIr_X77lupc31YXg,23964
-brainstate/transform/_autograd_test.py,sha256=epQ2z97fAp_dQ_CwWGZD7sgw-p9o9fGfSeOUAJiiDY0,38658
+brainstate/optim/_sgd_optimizer.py,sha256=JiK_AVGregL0wn8uHhRQvK9Qq7Qja7dEyLW6Aan7b70,45826
+brainstate/transform/__init__.py,sha256=my2X4ZW0uKZRfN82zyGEPizWNJ0fsSP2akvmkjn43ck,1458
+brainstate/transform/_autograd.py,sha256=Pj_YxpU52guaxQs1NcB6qDtXgkvaPcoJbuvIF8T-Wmk,23964
+brainstate/transform/_autograd_test.py,sha256=RWriMemIF9FVFUjQh4IHzLhT9LGyd1JXpjXfFZKHn10,38654
 brainstate/transform/_control.py,sha256=NWceTIuLlj2uGTdNcqBAXgnaLuChOGgAtIXtFn5vdLU,26837
 brainstate/transform/_controls_test.py,sha256=mPUa_qmXXVxDziAJrPWRBwsGnc3cHR9co08eJB_fJwA,7648
 brainstate/transform/_jit.py,sha256=sjQHFV8Tt75fpdl12jjPRDPT92_IZxBBJAG4gapdbNQ,11471
 brainstate/transform/_jit_error.py,sha256=lO_e5AdhkjozHjM10q0b57OaXbeZ9gQkVmZMN6VQVCw,4450
-brainstate/transform/_jit_test.py,sha256=lVXvScfXExhXwFi8jnvEY6stNVulZHCzriamajFqzrY,2891
-brainstate/transform/_make_jaxpr.py,sha256=MTeBpPO1thu5yDytWoJijySHV7-nWmUoBMC0RCbdzcY,29972
-brainstate/transform/_make_jaxpr_test.py,sha256=4nEwZv_ebgUZgV86vOJFO_qC69mw2F3rogViF2SC1Qs,3823
-brainstate/transform/_progress_bar.py,sha256=myrAkBcUfuVGFLVwFzeSe5vdg1z49ARKqTlccG92maA,3536
-brainstate-0.0.1.post20240612.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.0.1.post20240612.dist-info/METADATA,sha256=VRHXnO0TBRcoo_M4iFHsywjCJbhonpStzYSmRXkR_wM,4254
-brainstate-0.0.1.post20240612.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.0.1.post20240612.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.0.1.post20240612.dist-info/RECORD,,
+brainstate/transform/_jit_test.py,sha256=5ltT7izh_OS9dcHnRymmVhq01QomjwZGdA8XzwJRLb4,2868
+brainstate/transform/_make_jaxpr.py,sha256=q3OPy-1Gg0mVaB9pgSTWzzP8FSCAgquSjP-pDEw3Tpg,29970
+brainstate/transform/_make_jaxpr_test.py,sha256=K3vRUBroDTCCx0lnmhgHtgrlWvWglJO2f1K2phTvU70,3819
+brainstate/transform/_progress_bar.py,sha256=VGoRZPRBmB8ELNwLc6c7S8QhUUTvn0FY46IbBm9cuYM,3502
+brainstate-0.0.1.post20240622.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.0.1.post20240622.dist-info/METADATA,sha256=N-I84Xg5_9MT9p4LjmQjpk4GcPV5iDR3J8LDa9ppKnM,3814
+brainstate-0.0.1.post20240622.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+brainstate-0.0.1.post20240622.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.0.1.post20240622.dist-info/RECORD,,

brainstate/math/__init__.py

@@ -1,21 +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 ._einops import *
-from ._einops import __all__ as _einops_all
-from ._misc import *
-from ._misc import __all__ as _misc_all
-
-__all__ = _misc_all + _einops_all