brainstate 0.0.1.1.post20240802__py2.py3-none-any.whl → 0.0.2__py2.py3-none-any.whl

brainstate/__init__.py

@@ -17,7 +17,7 @@
 A ``State``-based Transformation System for Brain Dynamics Programming
 """
 
-__version__ = "0.0.1.1"
+__version__ = "0.0.2"
 
 from . import environ
 from . import functional

brainstate/_module.py

@@ -61,11 +61,9 @@ from ._state import State, StateDictManager, visible_state_dict
 from ._utils import set_module_as
 from .mixin import Mixin, Mode, DelayedInit, JointTypes, Batching, UpdateReturn
 from .transform import jit_error
+from .typing import Size, ArrayLike, PyTree
 from .util import unique_name, DictManager, get_unique_name
 
-Shape = Union[int, Sequence[int]]
-PyTree = Any
-ArrayLike = jax.typing.ArrayLike
 
 delay_identifier = '_*_delay_of_'
 _DELAY_ROTATE = 'rotation'
@@ -805,7 +803,7 @@ class Dynamics(ExtendedUpdateWithBA, ReceiveInputProj, UpdateReturn):
 
   def __init__(
       self,
-      size: Shape,
+      size: Size,
       keep_size: bool = False,
       name: Optional[str] = None,
       mode: Optional[Mode] = None,
@@ -1275,25 +1273,25 @@ class Delay(ExtendedUpdateWithBA, DelayedInit):
 
     if self.interp_method == _INTERP_LINEAR:  # "linear" interpolation
       # def _interp(target):
-        # if len(indices) > 0:
-        #   raise NotImplementedError('The slicing indices are not supported in the linear interpolation.')
-        # if self.delay_method == _DELAY_ROTATE:
-        #   i = environ.get(environ.I, desc='The time step index.')
-        #   _interp_fun = partial(jnp.interp, period=self.max_length)
-        #   for dim in range(1, target.ndim, 1):
-        #     _interp_fun = jax.vmap(_interp_fun, in_axes=(None, None, dim), out_axes=dim - 1)
-        #   di = i - jnp.arange(self.max_length)
-        #   delay_idx = jnp.asarray(di % self.max_length, dtype=jnp.int32)
-        #   return _interp_fun(float_time_step, delay_idx, target)
-        #
-        # elif self.delay_method == _DELAY_CONCAT:
-        #   _interp_fun = partial(jnp.interp, period=self.max_length)
-        #   for dim in range(1, target.ndim, 1):
-        #     _interp_fun = jax.vmap(_interp_fun, in_axes=(None, None, dim), out_axes=dim - 1)
-        #   return _interp_fun(float_time_step, jnp.arange(self.max_length), target)
-        #
-        # else:
-        #   raise ValueError(f'Unknown delay updating method "{self.delay_method}"')
+      # if len(indices) > 0:
+      #   raise NotImplementedError('The slicing indices are not supported in the linear interpolation.')
+      # if self.delay_method == _DELAY_ROTATE:
+      #   i = environ.get(environ.I, desc='The time step index.')
+      #   _interp_fun = partial(jnp.interp, period=self.max_length)
+      #   for dim in range(1, target.ndim, 1):
+      #     _interp_fun = jax.vmap(_interp_fun, in_axes=(None, None, dim), out_axes=dim - 1)
+      #   di = i - jnp.arange(self.max_length)
+      #   delay_idx = jnp.asarray(di % self.max_length, dtype=jnp.int32)
+      #   return _interp_fun(float_time_step, delay_idx, target)
+      #
+      # elif self.delay_method == _DELAY_CONCAT:
+      #   _interp_fun = partial(jnp.interp, period=self.max_length)
+      #   for dim in range(1, target.ndim, 1):
+      #     _interp_fun = jax.vmap(_interp_fun, in_axes=(None, None, dim), out_axes=dim - 1)
+      #   return _interp_fun(float_time_step, jnp.arange(self.max_length), target)
+      #
+      # else:
+      #   raise ValueError(f'Unknown delay updating method "{self.delay_method}"')
       # return jax.tree.map(_interp, self.history.value)
 
       data_at_t0 = self.retrieve_at_step(jnp.asarray(jnp.floor(float_time_step), dtype=jnp.int32), *indices)

brainstate/_random_for_unit.py

@@ -30,13 +30,14 @@ def uniform_for_unit(
     maxval: ArrayLike = 1.
 ) -> jax.Array | bu.Quantity:
   if isinstance(minval, bu.Quantity) and isinstance(maxval, bu.Quantity):
-    return bu.Quantity(jr.uniform(key, shape, dtype, minval.value, maxval.value), dim=minval.dim)
+    maxval = maxval.in_unit(minval.unit)
+    return bu.Quantity(jr.uniform(key, shape, dtype, minval.mantissa, maxval.mantissa), unit=minval.unit)
   elif isinstance(minval, bu.Quantity):
     assert minval.is_unitless, f'minval must be unitless when maxval is not a Quantity, got {minval}'
-    minval = minval.value
+    minval = minval.mantissa
   elif isinstance(maxval, bu.Quantity):
     assert maxval.is_unitless, f'maxval must be unitless when minval is not a Quantity, got {maxval}'
-    maxval = maxval.value
+    maxval = maxval.mantissa
   return jr.uniform(key, shape, dtype, minval, maxval)
 
 
@@ -47,5 +48,5 @@ def permutation_for_unit(
     independent: bool = False
 ) -> jax.Array | bu.Quantity:
   if isinstance(x, bu.Quantity):
-    return bu.Quantity(jr.permutation(key, x.value, axis, independent=independent), dim=x.dim)
+    return bu.Quantity(jr.permutation(key, x.mantissa, axis, independent=independent), unit=x.unit)
   return jr.permutation(key, x, axis, independent=independent)

brainstate/_state.py

@@ -22,11 +22,9 @@ import numpy as np
 from jax.api_util import shaped_abstractify
 from jax.extend import source_info_util
 
+from .typing import ArrayLike, PyTree
 from .util import DictManager
 
-PyTree = Any
-max_int = np.iinfo(np.int32)
-
 __all__ = [
   'State', 'ShortTermState', 'LongTermState', 'ParamState',
   'StateDictManager',
@@ -36,6 +34,7 @@ __all__ = [
 ]
 
 _pytree_registered_objects = set()
+max_int = np.iinfo(np.int32)
 
 
 def _register_pytree_cls(cls):
@@ -110,7 +109,7 @@ class State(object):
   __module__ = 'brainstate'
   __slots__ = ('_value', '_name', '_tree', '_level', '_source_info', '_check_tree')
 
-  def __init__(self, value: PyTree, name: Optional[str] = None):
+  def __init__(self, value: PyTree[ArrayLike], name: Optional[str] = None):
     if isinstance(value, State):
       value = value.value
     self._value = value
@@ -135,7 +134,7 @@ class State(object):
     self._name = name
 
   @property
-  def value(self) -> PyTree:
+  def value(self) -> PyTree[ArrayLike]:
     """
     The data and its value.
     """

brainstate/functional/_activations.py

@@ -25,8 +25,8 @@ from typing import Any, Union, Sequence
 import jax
 import jax.numpy as jnp
 from jax.scipy.special import logsumexp
-from jax.typing import ArrayLike
 
+from brainstate.typing import ArrayLike
 from .. import random
 
 __all__ = [

brainstate/functional/_others.py

@@ -16,12 +16,11 @@
 from __future__ import annotations
 
 from functools import partial
-from typing import Any
 
 import jax
 import jax.numpy as jnp
 
-PyTree = Any
+from brainstate.typing import PyTree
 
 __all__ = [
   'clip_grad_norm',

brainstate/init/_generic.py

@@ -22,8 +22,8 @@ import jax
 import numpy as np
 
 from brainstate._state import State
+from brainstate.typing import ArrayLike
 from ._base import to_size
-from ..typing import ArrayLike
 
 __all__ = [
   'param',

brainstate/init/_random_inits.py

@@ -22,8 +22,8 @@ import jax.numpy as jnp
 import numpy as np
 
 from brainstate import environ, random
+from brainstate.typing import ArrayLike
 from ._base import Initializer, to_size
-from ..typing import ArrayLike
 
 __all__ = [
   'Normal',

brainstate/mixin.py

@@ -15,14 +15,15 @@
 
 # -*- coding: utf-8 -*-
 
-from typing import Sequence, Optional, TypeVar, Any
+from typing import Sequence, Optional, TypeVar
 from typing import (_SpecialForm, _type_check, _remove_dups_flatten, _UnionGenericAlias)
 
-T = TypeVar('T')
-PyTree = Any
+from .typing import PyTree
 
+T = TypeVar('T')
 State = None
 
+
 __all__ = [
   'Mixin',
   'DelayedInit',

brainstate/random.py

@@ -490,14 +490,13 @@ class RandomState(State):
     upper = bu.math.asarray(upper, dtype=dtype)
     loc = bu.math.asarray(loc, dtype=dtype)
     scale = bu.math.asarray(scale, dtype=dtype)
-    bu.fail_for_dimension_mismatch(lower, upper)
-    bu.fail_for_dimension_mismatch(lower, loc)
-    bu.fail_for_dimension_mismatch(lower, scale)
-    dim = lower.dim if isinstance(lower, bu.Quantity) else bu.DIMENSIONLESS
-    lower = lower.value if isinstance(lower, bu.Quantity) else lower
-    upper = upper.value if isinstance(upper, bu.Quantity) else upper
-    loc = loc.value if isinstance(loc, bu.Quantity) else loc
-    scale = scale.value if isinstance(scale, bu.Quantity) else scale
+    unit = bu.get_unit(lower)
+    lower, upper, loc, scale = (
+      lower.mantissa if isinstance(lower, bu.Quantity) else lower,
+      bu.Quantity(upper).in_unit(unit).mantissa,
+      bu.Quantity(loc).in_unit(unit).mantissa,
+      bu.Quantity(scale).in_unit(unit).mantissa
+    )
 
     jit_error(
       bu.math.any(bu.math.logical_or(loc < lower - 2 * scale, loc > upper + 2 * scale)),
@@ -535,10 +534,12 @@ class RandomState(State):
     out = out * scale * sqrt2 + loc
 
     # Clamp to ensure it's in the proper range
-    out = jnp.clip(out,
-                   lax.nextafter(lax.stop_gradient(lower), np.array(np.inf, dtype=dtype)),
-                   lax.nextafter(lax.stop_gradient(upper), np.array(-np.inf, dtype=dtype)))
-    return out if dim == bu.DIMENSIONLESS else bu.Quantity(out, dim=dim)
+    out = jnp.clip(
+      out,
+      lax.nextafter(lax.stop_gradient(lower), np.array(np.inf, dtype=dtype)),
+      lax.nextafter(lax.stop_gradient(upper), np.array(-np.inf, dtype=dtype))
+    )
+    return out if unit.is_unitless else bu.Quantity(out, unit=unit)
 
   def _check_p(self, p):
     raise ValueError(f'Parameter p should be within [0, 1], but we got {p}')
@@ -555,30 +556,33 @@ class RandomState(State):
     r = jr.bernoulli(key, p=p, shape=_size2shape(size))
     return r
 
-  def lognormal(self,
-                mean=None,
-                sigma=None,
-                size: Optional[Size] = None,
-                key: Optional[SeedOrKey] = None,
-                dtype: DTypeLike = None):
+  def lognormal(
+      self,
+      mean=None,
+      sigma=None,
+      size: Optional[Size] = None,
+      key: Optional[SeedOrKey] = None,
+      dtype: DTypeLike = None
+  ):
     mean = _check_py_seq(mean)
     sigma = _check_py_seq(sigma)
     mean = bu.math.asarray(mean, dtype=dtype)
     sigma = bu.math.asarray(sigma, dtype=dtype)
-    bu.fail_for_dimension_mismatch(mean, sigma)
-    dim = mean.dim if isinstance(mean, bu.Quantity) else bu.DIMENSIONLESS
-    mean = mean.value if isinstance(mean, bu.Quantity) else mean
-    sigma = sigma.value if isinstance(sigma, bu.Quantity) else sigma
+    unit = mean.unit if isinstance(mean, bu.Quantity) else bu.Unit()
+    mean = mean.mantissa if isinstance(mean, bu.Quantity) else mean
+    sigma = sigma.in_unit(unit).mantissa if isinstance(sigma, bu.Quantity) else sigma
 
     if size is None:
-      size = jnp.broadcast_shapes(jnp.shape(mean),
-                                  jnp.shape(sigma))
+      size = jnp.broadcast_shapes(
+        jnp.shape(mean),
+        jnp.shape(sigma)
+      )
     key = self.split_key() if key is None else _formalize_key(key)
     dtype = dtype or environ.dftype()
     samples = jr.normal(key, shape=_size2shape(size), dtype=dtype)
     samples = _loc_scale(mean, sigma, samples)
     samples = jnp.exp(samples)
-    return samples if dim == bu.DIMENSIONLESS else bu.Quantity(samples, dim=dim)
+    return samples if unit.is_unitless else bu.Quantity(samples, unit=unit)
 
   def binomial(self,
                n,
@@ -678,10 +682,10 @@ class RandomState(State):
     cov = bu.math.asarray(_check_py_seq(cov), dtype=dtype)
     if isinstance(mean, bu.Quantity):
       assert isinstance(cov, bu.Quantity)
-      assert mean.dim ** 2 == cov.dim
-    mean = mean.value if isinstance(mean, bu.Quantity) else mean
-    cov = cov.value if isinstance(cov, bu.Quantity) else cov
-    dim = mean.dim if isinstance(mean, bu.Quantity) else bu.DIMENSIONLESS
+      assert mean.unit ** 2 == cov.unit
+    mean = mean.mantissa if isinstance(mean, bu.Quantity) else mean
+    cov = cov.mantissa if isinstance(cov, bu.Quantity) else cov
+    unit = mean.unit if isinstance(mean, bu.Quantity) else bu.Unit()
 
     key = self.split_key() if key is None else _formalize_key(key)
     if not jnp.ndim(mean) >= 1:
@@ -708,7 +712,7 @@ class RandomState(State):
       factor = jnp.linalg.cholesky(cov)
     normal_samples = jr.normal(key, size + mean.shape[-1:], dtype=dtype)
     r = mean + jnp.einsum('...ij,...j->...i', factor, normal_samples)
-    return r if dim == bu.DIMENSIONLESS else bu.Quantity(r, dim=dim)
+    return r if unit.is_unitless else bu.Quantity(r, unit=unit)
 
   def rayleigh(self,
                scale=1.0,
@@ -4809,10 +4813,9 @@ def _size2shape(size):
 
 
 def _check_shape(name, shape, *param_shapes):
-  shape = core.as_named_shape(shape)
   if param_shapes:
-    shape_ = lax.broadcast_shapes(shape.positional, *param_shapes)
-    if shape.positional != shape_:
+    shape_ = lax.broadcast_shapes(shape, *param_shapes)
+    if shape != shape_:
       msg = ("{} parameter shapes must be broadcast-compatible with shape "
              "argument, and the result of broadcasting the shapes must equal "
              "the shape argument, but got result {} for shape argument {}.")

brainstate/transform/_jit_error_test.py

@@ -16,8 +16,8 @@
 import unittest
 
 import jax
-import jaxlib.xla_extension
 import jax.numpy as jnp
+import jaxlib.xla_extension
 
 import brainstate as bst
 

brainstate/transform/_make_jaxpr.py

@@ -71,8 +71,8 @@ from jax.util import wraps
 
 from brainstate._state import State, StateTrace
 from brainstate._utils import set_module_as
+from brainstate.typing import PyTree
 
-PyTree = Any
 AxisName = Hashable
 
 __all__ = [

brainstate/typing.py

@@ -14,13 +14,16 @@
 # ==============================================================================
 
 
-from typing import Any, Sequence, Protocol, Union
+import functools as ft
+import typing
+from typing import Sequence, Protocol, Union, Any, Generic, TypeVar
 
 import brainunit as bu
 import jax
 import numpy as np
 
 __all__ = [
+  'PyTree',
   'Size',
   'Axes',
   'SeedOrKey',
@@ -29,6 +32,151 @@ __all__ = [
   'DTypeLike',
 ]
 
+_T = TypeVar("_T")
+
+
+class _FakePyTree(Generic[_T]):
+  pass
+
+
+_FakePyTree.__name__ = "PyTree"
+_FakePyTree.__qualname__ = "PyTree"
+_FakePyTree.__module__ = "builtins"
+
+
+class _MetaPyTree(type):
+  def __call__(self, *args, **kwargs):
+    raise RuntimeError("PyTree cannot be instantiated")
+
+  # Can't return a generic (e.g. _FakePyTree[item]) because generic aliases don't do
+  # the custom __instancecheck__ that we want.
+  # We can't add that __instancecheck__  via subclassing, e.g.
+  # type("PyTree", (Generic[_T],), {}), because dynamic subclassing of typeforms
+  # isn't allowed.
+  # Likewise we can't do types.new_class("PyTree", (Generic[_T],), {}) because that
+  # has __module__ "types", e.g. we get types.PyTree[int].
+  @ft.lru_cache(maxsize=None)
+  def __getitem__(cls, item):
+    if isinstance(item, tuple):
+      if len(item) == 2:
+
+        class X(PyTree):
+          leaftype = item[0]
+          structure = item[1].strip()
+
+        if not isinstance(X.structure, str):
+          raise ValueError(
+            "The structure annotation `struct` in "
+            "`brainstate.typing.PyTree[leaftype, struct]` must be be a string, "
+            f"e.g. `brainstate.typing.PyTree[leaftype, 'T']`. Got '{X.structure}'."
+          )
+        pieces = X.structure.split()
+        if len(pieces) == 0:
+          raise ValueError(
+            "The string `struct` in `brainstate.typing.PyTree[leaftype, struct]` "
+            "cannot be the empty string."
+          )
+        for piece_index, piece in enumerate(pieces):
+          if (piece_index == 0) or (piece_index == len(pieces) - 1):
+            if piece == "...":
+              continue
+          if not piece.isidentifier():
+            raise ValueError(
+              "The string `struct` in "
+              "`brainstate.typing.PyTree[leaftype, struct]` must be be a "
+              "whitespace-separated sequence of identifiers, e.g. "
+              "`brainstate.typing.PyTree[leaftype, 'T']` or "
+              "`brainstate.typing.PyTree[leaftype, 'foo bar']`.\n"
+              "(Here, 'identifier' is used in the same sense as in "
+              "regular Python, i.e. a valid variable name.)\n"
+              f"Got piece '{piece}' in overall structure '{X.structure}'."
+            )
+        name = str(_FakePyTree[item[0]])[:-1] + ', "' + item[1].strip() + '"]'
+      else:
+        raise ValueError(
+          "The subscript `foo` in `brainstate.typing.PyTree[foo]` must either be a "
+          "leaf type, e.g. `PyTree[int]`, or a 2-tuple of leaf and "
+          "structure, e.g. `PyTree[int, 'T']`. Received a tuple of length "
+          f"{len(item)}."
+        )
+    else:
+      name = str(_FakePyTree[item])
+
+      class X(PyTree):
+        leaftype = item
+        structure = None
+
+    X.__name__ = name
+    X.__qualname__ = name
+    if getattr(typing, "GENERATING_DOCUMENTATION", False):
+      X.__module__ = "builtins"
+    else:
+      X.__module__ = "brainstate.typing"
+    return X
+
+
+# Can't do `class PyTree(Generic[_T]): ...` because we need to override the
+# instancecheck for PyTree[foo], but subclassing
+# `type(Generic[int])`, i.e. `typing._GenericAlias` is disallowed.
+PyTree = _MetaPyTree("PyTree", (), {})
+if getattr(typing, "GENERATING_DOCUMENTATION", False):
+  PyTree.__module__ = "builtins"
+else:
+  PyTree.__module__ = "brainstate.typing"
+PyTree.__doc__ = """Represents a PyTree.
+
+Annotations of the following sorts are supported:
+```python
+a: PyTree
+b: PyTree[LeafType]
+c: PyTree[LeafType, "T"]
+d: PyTree[LeafType, "S T"]
+e: PyTree[LeafType, "... T"]
+f: PyTree[LeafType, "T ..."]
+```
+
+These correspond to:
+
+a. A plain `PyTree` can be used an annotation, in which case `PyTree` is simply a
+    suggestively-named alternative to `Any`.
+    ([By definition all types are PyTrees.](https://jax.readthedocs.io/en/latest/pytrees.html))
+
+b. `PyTree[LeafType]` denotes a PyTree all of whose leaves match `LeafType`. For
+    example, `PyTree[int]` or `PyTree[Union[str, Float32[Array, "b c"]]]`.
+
+c. A structure name can also be passed. In this case
+    `jax.tree_util.tree_structure(...)` will be called, and bound to the structure name.
+    This can be used to mark that multiple PyTrees all have the same structure:
+    ```python
+    def f(x: PyTree[int, "T"], y: PyTree[int, "T"]):
+        ...
+    ```
+
+d. A composite structure can be declared. In this case the variable must have a PyTree
+    structure each to the composition of multiple previously-bound PyTree structures.
+    For example:
+    ```python
+    def f(x: PyTree[int, "T"], y: PyTree[int, "S"], z: PyTree[int, "S T"]):
+        ...
+
+    x = (1, 2)
+    y = {"key": 3}
+    z = {"key": (4, 5)}  # structure is the composition of the structures of `y` and `z`
+    f(x, y, z)
+    ```
+    When performing runtime type-checking, all the individual pieces must have already
+    been bound to structures, otherwise the composite structure check will throw an error.
+
+e. A structure can begin with a `...`, to denote that the lower levels of the PyTree
+    must match the declared structure, but the upper levels can be arbitrary. As in the
+    previous case, all named pieces must already have been seen and their structures
+    bound.
+
+f. A structure can end with a `...`, to denote that the PyTree must be a prefix of the
+    declared structure, but the lower levels can be arbitrary. As in the previous two
+    cases, all named pieces must already have been seen and their structures bound.
+"""  # noqa: E501
+
 Size = Union[int, Sequence[int]]
 Axes = Union[int, Sequence[int]]
 SeedOrKey = Union[int, jax.Array, np.ndarray]
@@ -44,7 +192,7 @@ ArrayLike = Union[
   np.ndarray,  # NumPy array type
   np.bool_, np.number,  # NumPy scalar types
   bool, int, float, complex,  # Python scalar types
-  bu.Quantity,  # quantity
+  bu.Quantity,  # Quantity
 ]
 
 # --- Dtype --- #

{brainstate-0.0.1.1.post20240802.dist-info → brainstate-0.0.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.0.1.1.post20240802
+Version: 0.0.2
 Summary: A State-based Transformation System for Brain Dynamics Programming.
 Home-page: https://github.com/brainpy/brainstate
 Author: BDP

{brainstate-0.0.1.1.post20240802.dist-info → brainstate-0.0.2.dist-info}/RECORD

@@ -1,27 +1,27 @@
-brainstate/__init__.py,sha256=oxslZrm6wxtBQDqwJFb2BaAKZFmnp4d_esDkaeuGMWE,1410
-brainstate/_module.py,sha256=UjhfmY26VHQ-kF6U4l68AslGU6vhfD-dR7gF-4Io5ic,52520
+brainstate/__init__.py,sha256=zipNSih9Tyvi4-5cXqNPGsDF7VeestkLp-lcjJ4-dA0,1408
+brainstate/_module.py,sha256=YJDp9aD38wBa_lY6BojWjWV9LJ2aFMAMYh-KZe5a4eM,52443
 brainstate/_module_test.py,sha256=oQaoaZBTo1o3wHrMEJTInQCc7RdcVs1gcfQGvdSb1SI,7843
-brainstate/_random_for_unit.py,sha256=eW4NJkX27VCCNWUwAlyt2otkeEthGKOpUoX6XJ6i95Y,1946
-brainstate/_state.py,sha256=t4lEikvxTfeL2TW0chLUvsQuuRoJSO-iXylUydl1i7k,12057
+brainstate/_random_for_unit.py,sha256=1rHr7gfH_bYrJfpxbDhQUk_j00Yosx-GzyZCXrLxsd0,2007
+brainstate/_state.py,sha256=C0widCOj_ca6zfqh95jzFXf_G5vi0hJyuQ5GIqEqOUs,12102
 brainstate/_state_test.py,sha256=HDdipndRLhEHWEdTmyT1ayEBkbv6qJKykfCWKI6yJ_E,1253
 brainstate/_utils.py,sha256=RLorgGJkt2BhbX4C-ygd-PPG0wfcGCghjSP93sRvzqM,833
 brainstate/environ.py,sha256=LwRwnFaTbv8l7nHRIbSV46WzcN7pGLQFhT_xDUox2yA,10240
-brainstate/mixin.py,sha256=2f2toMUmgJIiovX1wi8OIBM8sbWH6s9Usa1ixL9J4tg,10747
+brainstate/mixin.py,sha256=OumTTSVyYSbtudjfS_MRThsBaeVJ_0JggeMClY7xtBA,10758
 brainstate/mixin_test.py,sha256=-Ej9oUOu8O1M4oy37SVMj7xNRYhHHyAHwrjS_aISayo,2923
-brainstate/random.py,sha256=UbXfC0nrxk5FOsld0rCxBp2bIaeQHH5bj-NWQTR8bbQ,188447
+brainstate/random.py,sha256=BqEBYVD9TGe8dSzp8U0suK0O4r6Ox59GCq0mwfUndVQ,188073
 brainstate/random_test.py,sha256=cCeuYvlZkCS2_RgG0vipZFNSHG8b-uJ7SXM9SZDCYQM,17866
 brainstate/surrogate.py,sha256=1kgbn82GSlpReIytIVl29yozk75gkdZv0gTBlixQ4C4,43798
-brainstate/typing.py,sha256=tZyXpopNtp1UvqbntnmlUP6OBTM_ywyV9QBQ6Gu3IdU,2232
+brainstate/typing.py,sha256=6BlkLSN5TiaNO49q8b0OYyzcuSxmdoG3noIJTbyhE3s,7895
 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/_activations.py,sha256=gvZ9E1-TsEUlyO7Om0eYzlM9DF-14_A32-gta1mjGo4,17798
 brainstate/functional/_normalization.py,sha256=IxE580waloZylZVXcpUUK4bWQdlE6oSPfafaKYfDkbg,2169
-brainstate/functional/_others.py,sha256=ifB-l82y7ZB632yLUJOEcpkRY-yOoiJ0mtDOxNilp4M,1711
+brainstate/functional/_others.py,sha256=1Epp75RkGYobMc2kHISZuS-_xnTFk3zHb1UHadwugCo,1711
 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=G5WHSQs3M9htREIZ_OXqj1ffSF_BdYdPTaMqe9AKj-k,7315
-brainstate/init/_random_inits.py,sha256=ycxH9WyKvPhRbk9PkamCLlc9aX5YokREBfSCbrbFQQ4,16021
+brainstate/init/_generic.py,sha256=LB7IQfswOG6X-q0QX5N8T5vZmxdygetsSBQ6iXlZ0oU,7324
+brainstate/init/_random_inits.py,sha256=LsfvKSX4wsR7Kh5jgKgdyXTCEEa5Nn_iYcp_2GgLQKY,16030
 brainstate/init/_regular_inits.py,sha256=u77aSM0BkK9VULFJQZ1lIEYA_sJJzEZBTEttBSJ79RI,3090
 brainstate/nn/__init__.py,sha256=YJHoI8cXKVRS8f2vUl3Zegp5wm0svMz3qo9JmQJiMQk,2162
 brainstate/nn/_base.py,sha256=lzbZpku3Q2arH6ZaAwjs6bhbV0RcFChxo2UcpnX5t84,8481
@@ -54,13 +54,13 @@ brainstate/transform/_control.py,sha256=0NFUGLIenqKuBhBiTmY0YgCrl2GI1ZbuWMW0DSOo
 brainstate/transform/_controls_test.py,sha256=mPUa_qmXXVxDziAJrPWRBwsGnc3cHR9co08eJB_fJwA,7648
 brainstate/transform/_jit.py,sha256=sjQHFV8Tt75fpdl12jjPRDPT92_IZxBBJAG4gapdbNQ,11471
 brainstate/transform/_jit_error.py,sha256=8rGRx8dtvmPWmHVOsfz30EUMXSix-m2PKM3Ni_9-_7I,4829
-brainstate/transform/_jit_error_test.py,sha256=MGFSCih6GVXTmLh0K-m8RX8N4x2_CAgtcGOIWfzol88,1862
+brainstate/transform/_jit_error_test.py,sha256=GAVGL0eNJ5Fu0lHABCGc-nLfa_0x0tw_VPfURB-nhLc,1862
 brainstate/transform/_jit_test.py,sha256=5ltT7izh_OS9dcHnRymmVhq01QomjwZGdA8XzwJRLb4,2868
-brainstate/transform/_make_jaxpr.py,sha256=vMPruKfp5Ugv8RL-9wGfQdSumLZdLtThZvv3sU9MDjE,30426
+brainstate/transform/_make_jaxpr.py,sha256=ZkrOZu4_0xcILuPUA3RFEkorJ-xbDuDtXorJI_qVThE,30450
 brainstate/transform/_make_jaxpr_test.py,sha256=K3vRUBroDTCCx0lnmhgHtgrlWvWglJO2f1K2phTvU70,3819
 brainstate/transform/_progress_bar.py,sha256=VGoRZPRBmB8ELNwLc6c7S8QhUUTvn0FY46IbBm9cuYM,3502
-brainstate-0.0.1.1.post20240802.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.0.1.1.post20240802.dist-info/METADATA,sha256=LYHl7AJ94js5MhJVghX5mJ1KB_VA4kWVc1bCmZ0O3GY,3807
-brainstate-0.0.1.1.post20240802.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.0.1.1.post20240802.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.0.1.1.post20240802.dist-info/RECORD,,
+brainstate-0.0.2.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.0.2.dist-info/METADATA,sha256=K6yiVOqGj3Qs_vKGgQmFXZtlu8cS4r7EZXl_iyCjwh0,3792
+brainstate-0.0.2.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+brainstate-0.0.2.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.0.2.dist-info/RECORD,,