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

brainstate/random/_rand_state.py

@@ -16,7 +16,6 @@
 # -*- coding: utf-8 -*-
 from __future__ import annotations
 
-from collections import namedtuple
 from functools import partial
 from operator import index
 from typing import Optional
@@ -37,6 +36,8 @@ from ._random_for_unit import uniform_for_unit, permutation_for_unit
 
 __all__ = ['RandomState', 'DEFAULT', ]
 
+use_prng_key = True
+
 
 class RandomState(State):
     """RandomState that track the random generator state. """
@@ -56,12 +57,15 @@ class RandomState(State):
             if seed_or_key is None:
                 seed_or_key = np.random.randint(0, 100000, 2, dtype=np.uint32)
         if isinstance(seed_or_key, int):
-            key = jr.PRNGKey(seed_or_key)
+            key = jr.PRNGKey(seed_or_key) if use_prng_key else jr.key(seed_or_key)
         else:
-            if len(seed_or_key) != 2 and seed_or_key.dtype != np.uint32:
-                raise ValueError('key must be an array with dtype uint32. '
-                                 f'But we got {seed_or_key}')
-            key = seed_or_key
+            if jnp.issubdtype(seed_or_key.dtype, jax.dtypes.prng_key):
+                key = seed_or_key
+            else:
+                if len(seed_or_key) != 2 and seed_or_key.dtype != np.uint32:
+                    raise ValueError('key must be an array with dtype uint32. '
+                                     f'But we got {seed_or_key}')
+                key = seed_or_key
         super().__init__(key)
 
         self._backup = None
@@ -70,6 +74,9 @@ class RandomState(State):
         return f'{self.__class__.__name__}({self.value})'
 
     def check_if_deleted(self):
+        if not use_prng_key and isinstance(self._value, np.ndarray):
+            self._value = jr.key(np.random.randint(0, 10000))
+
         if (
             isinstance(self._value, jax.Array) and
             not isinstance(self._value, jax.core.Tracer) and
@@ -111,12 +118,19 @@ class RandomState(State):
             if seed_or_key is None:
                 seed_or_key = np.random.randint(0, 100000, 2, dtype=np.uint32)
         if np.size(seed_or_key) == 1:
-            key = jr.PRNGKey(seed_or_key)
+            if isinstance(seed_or_key, int):
+                key = jr.PRNGKey(seed_or_key) if use_prng_key else jr.key(seed_or_key)
+            elif jnp.issubdtype(seed_or_key.dtype, jax.dtypes.prng_key):
+                key = seed_or_key
+            elif isinstance(seed_or_key, (jnp.ndarray, np.ndarray)) and jnp.issubdtype(seed_or_key.dtype, jnp.integer):
+                key = jr.PRNGKey(seed_or_key) if use_prng_key else jr.key(seed_or_key)
+            else:
+                raise ValueError(f'Invalid seed_or_key: {seed_or_key}')
         else:
-            if len(seed_or_key) != 2 and seed_or_key.dtype != np.uint32:
-                raise ValueError('key must be an array with dtype uint32. '
-                                 f'But we got {seed_or_key}')
-            key = seed_or_key
+            if len(seed_or_key) == 2 and seed_or_key.dtype == np.uint32:
+                key = seed_or_key
+            else:
+                raise ValueError(f'Invalid seed_or_key: {seed_or_key}')
         self.value = key
 
     def split_key(self, n: Optional[int] = None, backup: bool = False) -> SeedOrKey:
@@ -560,15 +574,15 @@ class RandomState(State):
         )
         return out if unit.is_unitless else u.Quantity(out, unit=unit)
 
-    def _check_p(self, p):
-        raise ValueError(f'Parameter p should be within [0, 1], but we got {p}')
+    def _check_p(self, *args, **kwargs):
+        raise ValueError('Parameter p should be within [0, 1], but we got {p}')
 
     def bernoulli(self,
                   p,
                   size: Optional[Size] = None,
                   key: Optional[SeedOrKey] = None):
         p = _check_py_seq(p)
-        jit_error_if(jnp.any(jnp.logical_and(p < 0, p > 1)), self._check_p, p)
+        jit_error_if(jnp.any(jnp.logical_or(p < 0, p > 1)), self._check_p, p=p)
         if size is None:
             size = jnp.shape(p)
         key = self.split_key() if key is None else _formalize_key(key)
@@ -603,19 +617,27 @@ class RandomState(State):
         samples = jnp.exp(samples)
         return samples if unit.is_unitless else u.Quantity(samples, unit=unit)
 
-    def binomial(self,
-                 n,
-                 p,
-                 size: Optional[Size] = None,
-                 key: Optional[SeedOrKey] = None,
-                 dtype: DTypeLike = None):
+    def binomial(
+        self,
+        n,
+        p,
+        size: Optional[Size] = None,
+        key: Optional[SeedOrKey] = None,
+        dtype: DTypeLike = None,
+        check_valid: bool = True,
+    ):
         n = _check_py_seq(n)
         p = _check_py_seq(p)
-        jit_error_if(jnp.any(jnp.logical_and(p < 0, p > 1)), self._check_p, p)
+        if check_valid:
+            jit_error_if(
+                jnp.any(jnp.logical_or(p < 0, p > 1)),
+                'Parameter p should be within [0, 1], but we got {p}',
+                p=p
+            )
         if size is None:
             size = jnp.broadcast_shapes(jnp.shape(n), jnp.shape(p))
         key = self.split_key() if key is None else _formalize_key(key)
-        r = _binomial(key, p, n, shape=_size2shape(size))
+        r = jr.binomial(key, n, p, shape=_size2shape(size))
         dtype = dtype or environ.ditype()
         return jnp.asarray(r, dtype=dtype)
 
@@ -1142,8 +1164,13 @@ DEFAULT = RandomState(np.random.randint(0, 10000, size=2, dtype=np.uint32))
 
 def _formalize_key(key):
     if isinstance(key, int):
-        return jr.PRNGKey(key)
+        return jr.PRNGKey(key) if use_prng_key else jr.key(key)
     elif isinstance(key, (jax.Array, np.ndarray)):
+        if jnp.issubdtype(key.dtype, jax.dtypes.prng_key):
+            return key
+        if key.size == 1 and jnp.issubdtype(key.dtype, jnp.integer):
+            return jr.PRNGKey(key) if use_prng_key else jr.key(key)
+
         if key.dtype != jnp.uint32:
             raise TypeError('key must be a int or an array with two uint32.')
         if key.size != 2:
@@ -1216,162 +1243,6 @@ def _const(example, val):
     return np.array(val, dtype)
 
 
-_tr_params = namedtuple(
-    "tr_params", ["c", "b", "a", "alpha", "u_r", "v_r", "m", "log_p", "log1_p", "log_h"]
-)
-
-
-def _get_tr_params(n, p):
-    # See Table 1. Additionally, we pre-compute log(p), log1(-p) and the
-    # constant terms, that depend only on (n, p, m) in log(f(k)) (bottom of page 5).
-    mu = n * p
-    spq = jnp.sqrt(mu * (1 - p))
-    c = mu + 0.5
-    b = 1.15 + 2.53 * spq
-    a = -0.0873 + 0.0248 * b + 0.01 * p
-    alpha = (2.83 + 5.1 / b) * spq
-    u_r = 0.43
-    v_r = 0.92 - 4.2 / b
-    m = jnp.floor((n + 1) * p).astype(n.dtype)
-    log_p = jnp.log(p)
-    log1_p = jnp.log1p(-p)
-    log_h = ((m + 0.5) * (jnp.log((m + 1.0) / (n - m + 1.0)) + log1_p - log_p) +
-             _stirling_approx_tail(m) + _stirling_approx_tail(n - m))
-    return _tr_params(c, b, a, alpha, u_r, v_r, m, log_p, log1_p, log_h)
-
-
-def _stirling_approx_tail(k):
-    precomputed = jnp.array([0.08106146679532726,
-                             0.04134069595540929,
-                             0.02767792568499834,
-                             0.02079067210376509,
-                             0.01664469118982119,
-                             0.01387612882307075,
-                             0.01189670994589177,
-                             0.01041126526197209,
-                             0.009255462182712733,
-                             0.008330563433362871],
-                            dtype=environ.dftype())
-    kp1 = k + 1
-    kp1sq = (k + 1) ** 2
-    return jnp.where(k < 10,
-                     precomputed[k],
-                     (1.0 / 12 - (1.0 / 360 - (1.0 / 1260) / kp1sq) / kp1sq) / kp1)
-
-
-def _binomial_btrs(key, p, n):
-    """
-    Based on the transformed rejection sampling algorithm (BTRS) from the
-    following reference:
-
-    Hormann, "The Generation of Binonmial Random Variates"
-    (https://core.ac.uk/download/pdf/11007254.pdf)
-    """
-
-    def _btrs_body_fn(val):
-        _, key, _, _ = val
-        key, key_u, key_v = jr.split(key, 3)
-        u = jr.uniform(key_u)
-        v = jr.uniform(key_v)
-        u = u - 0.5
-        k = jnp.floor(
-            (2 * tr_params.a / (0.5 - jnp.abs(u)) + tr_params.b) * u + tr_params.c
-        ).astype(n.dtype)
-        return k, key, u, v
-
-    def _btrs_cond_fn(val):
-        def accept_fn(k, u, v):
-            # See acceptance condition in Step 3. (Page 3) of TRS algorithm
-            # v <= f(k) * g_grad(u) / alpha
-
-            m = tr_params.m
-            log_p = tr_params.log_p
-            log1_p = tr_params.log1_p
-            # See: formula for log(f(k)) at bottom of Page 5.
-            log_f = (
-                (n + 1.0) * jnp.log((n - m + 1.0) / (n - k + 1.0))
-                + (k + 0.5) * (jnp.log((n - k + 1.0) / (k + 1.0)) + log_p - log1_p)
-                + (_stirling_approx_tail(k) - _stirling_approx_tail(n - k))
-                + tr_params.log_h
-            )
-            g = (tr_params.a / (0.5 - jnp.abs(u)) ** 2) + tr_params.b
-            return jnp.log((v * tr_params.alpha) / g) <= log_f
-
-        k, key, u, v = val
-        early_accept = (jnp.abs(u) <= tr_params.u_r) & (v <= tr_params.v_r)
-        early_reject = (k < 0) | (k > n)
-        return lax.cond(
-            early_accept | early_reject,
-            (),
-            lambda _: ~early_accept,
-            (k, u, v),
-            lambda x: ~accept_fn(*x),
-        )
-
-    tr_params = _get_tr_params(n, p)
-    ret = lax.while_loop(
-        _btrs_cond_fn, _btrs_body_fn, (-1, key, 1.0, 1.0)
-    )  # use k=-1 initially so that cond_fn returns True
-    return ret[0]
-
-
-def _binomial_inversion(key, p, n):
-    def _binom_inv_body_fn(val):
-        i, key, geom_acc = val
-        key, key_u = jr.split(key)
-        u = jr.uniform(key_u)
-        geom = jnp.floor(jnp.log1p(-u) / log1_p) + 1
-        geom_acc = geom_acc + geom
-        return i + 1, key, geom_acc
-
-    def _binom_inv_cond_fn(val):
-        i, _, geom_acc = val
-        return geom_acc <= n
-
-    log1_p = jnp.log1p(-p)
-    ret = lax.while_loop(_binom_inv_cond_fn, _binom_inv_body_fn, (-1, key, 0.0))
-    return ret[0]
-
-
-def _binomial_dispatch(key, p, n):
-    def dispatch(key, p, n):
-        is_le_mid = p <= 0.5
-        pq = jnp.where(is_le_mid, p, 1 - p)
-        mu = n * pq
-        k = lax.cond(
-            mu < 10,
-            (key, pq, n),
-            lambda x: _binomial_inversion(*x),
-            (key, pq, n),
-            lambda x: _binomial_btrs(*x),
-        )
-        return jnp.where(is_le_mid, k, n - k)
-
-    # Return 0 for nan `p` or negative `n`, since nan values are not allowed for integer types
-    cond0 = jnp.isfinite(p) & (n > 0) & (p > 0)
-    return lax.cond(
-        cond0 & (p < 1),
-        (key, p, n),
-        lambda x: dispatch(*x),
-        (),
-        lambda _: jnp.where(cond0, n, 0),
-    )
-
-
-@partial(jit, static_argnums=(3,))
-def _binomial(key, p, n, shape):
-    shape = shape or lax.broadcast_shapes(jnp.shape(p), jnp.shape(n))
-    # reshape to map over axis 0
-    p = jnp.reshape(jnp.broadcast_to(p, shape), -1)
-    n = jnp.reshape(jnp.broadcast_to(n, shape), -1)
-    key = jr.split(key, jnp.size(p))
-    if jax.default_backend() == "cpu":
-        ret = lax.map(lambda x: _binomial_dispatch(*x), (key, p, n))
-    else:
-        ret = vmap(lambda *x: _binomial_dispatch(*x))(key, p, n)
-    return jnp.reshape(ret, shape)
-
-
 @partial(jit, static_argnums=(2,))
 def _categorical(key, p, shape):
     # this implementation is fast when event shape is small, and slow otherwise

brainstate/surrogate.py

@@ -19,9 +19,13 @@ from __future__ import annotations
 import jax
 import jax.numpy as jnp
 import jax.scipy as sci
-from jax.core import Primitive
 from jax.interpreters import batching, ad, mlir
 
+if jax.__version_info__ < (0, 4, 38):
+    from jax.core import Primitive
+else:
+    from jax.extend.core import Primitive
+
 __all__ = [
     'Surrogate',
     'Sigmoid',

brainstate/util/__init__.py

@@ -27,8 +27,6 @@ from ._scaling import *
 from ._scaling import __all__ as _mem_scale_all
 from ._struct import *
 from ._struct import __all__ as _struct_all
-from ._visualization import *
-from ._visualization import __all__ as _visualization_all
 
 __all__ = (
     _others_all
@@ -38,7 +36,6 @@ __all__ = (
     + _struct_all
     + _error_all
     + _mapping_all
-    + _visualization_all
 )
 del (
     _others_all,
@@ -48,5 +45,4 @@ del (
     _struct_all,
     _error_all,
     _mapping_all,
-    _visualization_all,
 )

brainstate/util/_caller.py

@@ -1,7 +1,7 @@
 # The file is adapted from the Flax library (https://github.com/google/flax).
 # The credit should go to the Flax authors.
 #
-# Copyright 2024 The Flax Authors & 2024 BDP Ecosystem.
+# Copyright 2024 The Flax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

brainstate/util/_dict.py

@@ -1,4 +1,7 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+# The file is adapted from the Flax library (https://github.com/google/flax).
+# The credit should go to the Flax authors.
+#
+# Copyright 2024 The Flax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

brainstate/util/_filter.py

@@ -1,7 +1,7 @@
 # The file is adapted from the Flax library (https://github.com/google/flax).
 # The credit should go to the Flax authors.
 #
-# Copyright 2024 The Flax Authors & 2024 BDP Ecosystem.
+# Copyright 2024 The Flax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

brainstate/util/_pretty_repr.py

@@ -1,7 +1,7 @@
 # The file is adapted from the Flax library (https://github.com/google/flax).
 # The credit should go to the Flax authors.
 #
-# Copyright 2024 The Flax Authors & 2024 BDP Ecosystem.
+# Copyright 2024 The Flax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

brainstate/util/_struct.py

@@ -1,7 +1,7 @@
 # The file is adapted from the Flax library (https://github.com/google/flax).
 # The credit should go to the Flax authors.
 #
-# Copyright 2024 The Flax Authors & 2024 BDP Ecosystem.
+# Copyright 2024 The Flax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.1.0.post20250105
+Version: 0.1.0.post20250120
 Summary: A ``State``-based Transformation System for Program Compilation and Augmentation.
 Home-page: https://github.com/chaobrain/brainstate
 Author: BrainState Developers
@@ -19,6 +19,7 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Programming Language :: Python
 Classifier: Topic :: Scientific/Engineering :: Bio-Informatics

{brainstate-0.1.0.post20250105.dist-info → brainstate-0.1.0.post20250120.dist-info}/RECORD

@@ -1,46 +1,45 @@
 brainstate/__init__.py,sha256=A-QKdOvSalsCMxgk80Iz6_xMiUin6con6JaONHfciSY,1526
-brainstate/_state.py,sha256=4aDpLyHGr1VlPXeLSfM3USQG5K4o7orF7IlaBdYrtfE,29098
-brainstate/_state_test.py,sha256=1boTp1w8DiCFLsPwNtlLrlIqGRpkasAmLid5bv2fgP4,2223
+brainstate/_state.py,sha256=GZ46liHZSHbAHQEuELvOeoJ27P9xiZDz06G2AASjAjA,29142
+brainstate/_state_test.py,sha256=rJUFRSXEqrrl4qANRewY9mnDlzSbtHwBIGeZ0ku-8Dg,1650
 brainstate/_utils.py,sha256=uJ6WWKq3yb05ZdktCQGLWOXsOJveL1H9pR7eev70Jes,1693
 brainstate/environ.py,sha256=PZnVFWPioUBuWmwCO8wwCKrHQfP3BR-5lYPRl5i5GDA,17698
 brainstate/environ_test.py,sha256=QD6sPCKNtqemVCGwkdImjMazatrvvLr6YeAVcfUnVVY,2045
 brainstate/mixin.py,sha256=g7uVUwZphZWsNs9pb48ozG2cDGaj0hs0g3lq8tDk-Sg,11310
 brainstate/mixin_test.py,sha256=Oq_0fwC9vpXDN4t4dTBhWzLdFDNlcYsrcip14F1yECI,3079
-brainstate/surrogate.py,sha256=YaY6RJ6kzpuPXWFjaWsxWt2MzJfdm5v_jeOR8V_jPoU,48369
+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/_autograd.py,sha256=o9ivoEY7BmtdM1XmzdMmeRXpj6Tvn5xNB8LSGp2HKC8,25238
 brainstate/augment/_autograd_test.py,sha256=S2eEgrwTzdSi3u2nKE3u37WSThosLwx1WCP9ptJAGKo,44060
-brainstate/augment/_eval_shape.py,sha256=dGlRVHOAZ9LSRZsFi1erxgEWHrnhBO3Kq3WW11-Hvng,3819
-brainstate/augment/_eval_shape_test.py,sha256=1nnxbU7hPRbZPQWNWbQ518pw-H7FGDKKnQpZGBY9uRI,1390
-brainstate/augment/_mapping.py,sha256=cpxzVGCEYnP5jPqrowYoPXciw_-QR2F3wggrRj1OCPc,21850
-brainstate/augment/_mapping_test.py,sha256=TEAecjZmTSDCfxARgrzcDJ2dW1Yz_sCITmFiA9FGrhk,9455
+brainstate/augment/_eval_shape.py,sha256=ObCgsZ704kLduB1dbjJZh5nVQYEkLR5ebK74V5NV42k,3892
+brainstate/augment/_eval_shape_test.py,sha256=LFOJx7CWltmRLXdGY175UebLwtEMz2CzJ_gLqMZsJTw,1393
+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/_ad_checkpoint.py,sha256=5zJ1ENeTU4FzRY_uNpr85NhKfuicMMjcIbhu6-bSM4k,9451
+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
 brainstate/compile/_conditions_test.py,sha256=s9LF6h9LvigvgxUIugTqvgCHBIU8TXS1Ar1OlIxXfrw,8389
 brainstate/compile/_error_if.py,sha256=TFvhqITKkRO9m30GdlUP4eEjJvLWQUhjkujXO9zvrWs,2689
 brainstate/compile/_error_if_test.py,sha256=OdJG483IIdOrCHxtHd49OHfOxCSnSkk7GdAUOzSt8bE,2044
-brainstate/compile/_jit.py,sha256=3WBXNTALWPYC9rQH0TPH6w4bjG0BpnZt3RAzUQF5kkc,14045
+brainstate/compile/_jit.py,sha256=itAWENKfJvnlaWl_uSy8lHTK8K1in89F_ZXXwp-EGRM,13944
 brainstate/compile/_jit_test.py,sha256=zD7kck9SQJGmUDolh9P4luKwQ21fBGje1Z4STTEXIuA,4135
-brainstate/compile/_loop_collect_return.py,sha256=XwMnKkMH0xTWB1f6GE4NQNK1R2GXTXCiVgulpkdIpc4,23308
+brainstate/compile/_loop_collect_return.py,sha256=TrKBZhtQecTtuiVz_HOeyepde-znzjlyk0Te53-AvOE,23492
 brainstate/compile/_loop_collect_return_test.py,sha256=bA-_11E8A_0jR5umEO3e409y7bb5QYDTgSL-SBaX7kQ,1802
-brainstate/compile/_loop_no_collection.py,sha256=0i31gdQ7sI-d6pvnh08ttUUwdAtpx4uoYhGuf_CyL9s,7343
+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=S5O9KUB3bsxoKcfptlV0MRfKA__Ija37WxkakIRL3z0,33010
+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=eInZPjiqzYE6PWxl_or_lBthDNDO0Ov60Uz0DbuBbZQ,4620
-brainstate/compile/_unvmap.py,sha256=0i-NvCLDAUe-effJIIEPVsK4WTPbCDBTgw6AqRvq7mE,4163
+brainstate/compile/_progress_bar.py,sha256=0oVlZ4kW_ZMciJjOR_ebj3PNe_XkCMkoQpv-HUUdoF0,5554
+brainstate/compile/_unvmap.py,sha256=EY4rbqCzzPOiaRwpWTiyBwb5dVkYFnacHhBZUZObxPI,4255
 brainstate/compile/_util.py,sha256=aCvkTV--g4NsqcodTdBAISt4EwgezCbKzNUV58n-Q_Y,6304
-brainstate/event/__init__.py,sha256=W0ZxbgrcFuYhWTl-GZ0UDoMGfsWmesvG4J_LbTBkex8,937
-brainstate/event/_csr.py,sha256=QDccbgXUklE2iq1w6WdyaFspXY1165uA9UlPltX16OU,30365
+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_mv_benchmark.py,sha256=xrj2DSWzw0pUHAE1jRBeSRhMW7ogXvDHEdeaZGioNE4,702
-brainstate/event/_csr_mv_test.py,sha256=WQfAvp_3UeCUGAZjr3_aqQvrB-eYZcFEN4v1PBe9fUQ,4012
-brainstate/event/_csr_test.py,sha256=v59rnwTy8jrvqjdGzN75kvLg0wLBmRbthaVRKY2f0Uw,2945
+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
@@ -48,7 +47,7 @@ brainstate/event/_linear_mv.py,sha256=O5qbY31GNV1qEDrZ5kvPbA8Ae-bY5JpUgGtqDFNAeV
 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=f0OrO6CjsJOUNUCRPpHIRmsb_wgNEym0xBl1tcz8ij4,11016
+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
@@ -56,14 +55,11 @@ brainstate/functional/_activations_test.py,sha256=T___RlSrIfXwlkw8dg5A9EZMTZGDzv
 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=ZUD-gY9txkr0yQ7VRqRSptzGMX3J_ZZ_VsJkxp8EfpY,1334
-brainstate/graph/_graph_context.py,sha256=J3WmCPrNyYiv-i75QSWQQdg0qZ6jmRy3gIzz3SuDBLI,15429
-brainstate/graph/_graph_context_test.py,sha256=IYpjqbXwSFF65XL0ZbdPeC1jYyEHLpQVrhuFeJXH4GM,2409
-brainstate/graph/_graph_convert.py,sha256=llSREtGQrIggkD0wmxUbYKuSveLW4ihDZME6Ab-mRTQ,9147
-brainstate/graph/_graph_node.py,sha256=mmZ0jhZev8ReNJhVLgWqYJEedEDtJHxhwxRv4ytQVNo,9268
+brainstate/graph/__init__.py,sha256=fyvQMlAUY3QYTzvDzz5TDoWS2XQwZ6P3ic6BtysZyHM,1026
+brainstate/graph/_graph_node.py,sha256=swAokZLKswSTaq2WEhyLIs38sy_67C6maHI6T3e1hvY,8339
 brainstate/graph/_graph_node_test.py,sha256=BFGfdzZFDHI0XK7hHotSVWKt3em1taGvn8FHF9NCXx8,2702
-brainstate/graph/_graph_operation.py,sha256=xM-pvYMLHb8e2lAmDVgTTThAf5gd4hHymfrqOJwpHeo,64132
-brainstate/graph/_graph_operation_test.py,sha256=ADyyuMk2xidEkkFNpGvUbvEtRmUj-tqOI4cF3eRuakM,24678
+brainstate/graph/_graph_operation.py,sha256=cIwGo3ICgtce2fmdn917r81evMFjJIKeW9doaQK4DD8,64111
+brainstate/graph/_graph_operation_test.py,sha256=zjvpKjQAFWtw8YZuqOk_jmlZNb_-E8oPyNx57dyc8jI,18556
 brainstate/init/__init__.py,sha256=R1dHgub47o-WJM9QkFLc7x_Q7GsyaKKDtrRHTFPpC5g,1097
 brainstate/init/_base.py,sha256=B_NLS9aKNrvuj5NAlSgBbQTVev7IRvzcx8vH0J-Gq2w,1671
 brainstate/init/_generic.py,sha256=sGOvd_atpxLWqqZKobTfAiMiYRnDC19PBNHdQy_igFM,8028
@@ -83,7 +79,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=pkcAVt_o5kQF_BGCTZZ-NUQpHgjlFHHPwtYC0fJkAA0,9099
+brainstate/nn/_dyn_impl/_inputs.py,sha256=UNoGxKIKXwPnhelljDowqAWlV6ds7aBBkEbvdy2oDI4,11302
 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
@@ -102,10 +98,10 @@ brainstate/nn/_elementwise/_dropout_test.py,sha256=ZzNvjFf46NpKWGBIcT6O0lKOBGpxO
 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
-brainstate/nn/_interaction/_conv.py,sha256=LgWO4TeKRru07UEUga3YX6xog6WHtOvKdKtgxGnHUvw,18512
+brainstate/nn/_interaction/_conv.py,sha256=lwyxTVsJVPiKlZcgB6iqE64aX7AOJzplDSj4y6-m18o,18592
 brainstate/nn/_interaction/_conv_test.py,sha256=fHXRFYnDghFiKre63RqMwIE_gbPKdK34UPhKOz-J3qU,8695
 brainstate/nn/_interaction/_embedding.py,sha256=iK0I1ExKWFa_QzV9UDGj32Ljsmdr1g_LlAtMcusebxU,2187
-brainstate/nn/_interaction/_linear.py,sha256=0G3OOmbgBwB22JAb2AWr4btvBOrFHBi4xWYKkd0SpOk,12148
+brainstate/nn/_interaction/_linear.py,sha256=EnkOk1oE79rvRIjU6HBllxUpVOEcQQCj4vtavo9AJjI,14767
 brainstate/nn/_interaction/_linear_test.py,sha256=QfCR8SBBed9OnSY-AmQ0kJqoggDA3Xem0dRJ0BusxLU,2872
 brainstate/nn/_interaction/_normalizations.py,sha256=7YDzkmO_iqd70fH_wawb60Bu8eGOdvZq23emP-b68Hc,37440
 brainstate/nn/_interaction/_normalizations_test.py,sha256=2p1Jf8nA999VYGWbvOZfKYlKk6UmL0vaEB76xkXxkXw,2438
@@ -119,26 +115,24 @@ brainstate/optim/_optax_optimizer.py,sha256=SuXV_xUBfhOw1_C2J5TIpy3dXDtI9VJFaSML
 brainstate/optim/_optax_optimizer_test.py,sha256=DAomE8Eu3dn4gh1S3EZ_u4pW4rhcl16vWPbnDcN3Rs4,1762
 brainstate/optim/_sgd_optimizer.py,sha256=NVKYhGcw2D1ksNWUIXZcj-74LUaan8XL3EERk-EHMRI,46008
 brainstate/random/__init__.py,sha256=c5q-RC3grRIjx-HBb2IhKZpi_xzbFmUUxzRAzqfREic,1045
-brainstate/random/_rand_funs.py,sha256=y3K39RZLoMAPv3mllnZNLGYTI30EHg6fFzlxAWOzvt0,137597
+brainstate/random/_rand_funs.py,sha256=WaelvEpeQb6Vuqt4eNgsAtd7GI8BqgEdVYbXgtCOd54,137682
 brainstate/random/_rand_funs_test.py,sha256=abO5lSoPBgBcg6ecFE1qnCg98__QGa68GSYC5pQW5QI,19438
-brainstate/random/_rand_seed.py,sha256=dfq-_vb-4YEEWL3Bkcm_VaVkxU2bkuOsIs3NlZy4BeE,5480
+brainstate/random/_rand_seed.py,sha256=MHA9znbdJW9ujx73onDRrAOI684_0FmGfqczBsSXYQg,5985
 brainstate/random/_rand_seed_test.py,sha256=Qibcs-ZqCvj1LuucmQ8H00B_HBNhf2f6un0aUdNZNTw,1518
-brainstate/random/_rand_state.py,sha256=EDTg4QCM1HxaLd_2F31mu-qTlQR5NFPDiKvI-QhgnOQ,59444
+brainstate/random/_rand_state.py,sha256=nuoQ8GU1MfJPRNN-ZmRQsggVjoyPhaEdZmwM7_4-Q3c,55206
 brainstate/random/_random_for_unit.py,sha256=kGp4EUX19MXJ9Govoivbg8N0bddqOldKEI2h_TbdONY,2057
-brainstate/util/__init__.py,sha256=EpAGQiukTUIVJWCfp106-aU1-jo7cIsCBBlvweXmwYY,1643
-brainstate/util/_caller.py,sha256=88-vWwIjOvNBIHkZSrBKfImeCJbyIPJqZ16pceC83yQ,2803
-brainstate/util/_dict.py,sha256=jv4QuHGNY3mgvmyS_2j6-HSex1lwcGlIjgaCACYpCtg,26077
+brainstate/util/__init__.py,sha256=-FWEuSKXG3mWxYphGFAy3UEuVe39lFs1GruluzdXDoI,1502
+brainstate/util/_caller.py,sha256=T3bzu7-09r-6EOrU6Muca_aMXSQua_X2lXjEqb-w39w,2782
+brainstate/util/_dict.py,sha256=Yapug-_RZQYjvd8cZ3v90_MX7rUYJDBzBnZJT6a0NXY,26178
 brainstate/util/_dict_test.py,sha256=Dn0TdjX6wLBXaTD4jfYTu6cKfFHwKSxi4_3bX7kB_IA,5621
 brainstate/util/_error.py,sha256=eyZ8PGFixqe2K5OEfjSDzI-2tU0ieYQoUpBP7yStlPQ,878
-brainstate/util/_filter.py,sha256=Jy_lUu1EC4HWgbJFkY6nhsTRc3xqONwZq36qopl7wRQ,4991
+brainstate/util/_filter.py,sha256=1-bvFHdjeehvXeHTrCEp8xr25lopKe8d3XZGCNegq0s,4970
 brainstate/util/_others.py,sha256=jsPZwP-v_5HRV-LB5F0NUsiqr04y8bmGIsu_JMyVcbQ,14762
-brainstate/util/_pretty_repr.py,sha256=NYEBCo2iz9Potx-IR7uZZzt2aLQW_94vH79fGusiC2A,5737
+brainstate/util/_pretty_repr.py,sha256=bDpU4gbkS4B8cXBkiN8kBQNmruxiJzDRF-eIqzyeYnM,5716
 brainstate/util/_scaling.py,sha256=pc_eM_SZVwkY65I4tJh1ODiHNCoEhsfFXl2zBK0PLAg,7562
-brainstate/util/_struct.py,sha256=0exv0oOiSt1hmx20Y4J2-pCGtCTx13WcAlEYSBkyung,17640
-brainstate/util/_tracers.py,sha256=0r5T4nhxMzI79NtqroqitsdMT4YfpgV5RdYJLS5uJ0w,2285
-brainstate/util/_visualization.py,sha256=n4ZVz10z7VBqA0cKO6vyHwEMprWJgPeEqtITzDMai2Y,1519
-brainstate-0.1.0.post20250105.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.1.0.post20250105.dist-info/METADATA,sha256=Xec1GNBlHcignyvym-EHzU-JIOUuo3T-IUU2LoCO0sk,3533
-brainstate-0.1.0.post20250105.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.1.0.post20250105.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.1.0.post20250105.dist-info/RECORD,,
+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,,