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

brainstate/__init__.py

@@ -17,7 +17,7 @@
 A ``State``-based Transformation System for Program Compilation and Augmentation
 """
 
-__version__ = "0.1.4"
+__version__ = "0.1.5"
 
 from . import augment
 from . import compile

brainstate/_state.py

@@ -50,6 +50,7 @@ __all__ = [
     'LongTermState',
     'HiddenState',
     'ParamState',
+    'BatchState',
     'TreefyState',
     'FakeState',
 

brainstate/augment/_mapping.py

@@ -185,10 +185,10 @@ def _compile_stateful_function(
     if isinstance(in_axes, int):
         args = jax.tree.map(lambda x: _remove_axis(x, in_axes), args)
     elif isinstance(in_axes, tuple):
-        args = tuple(
-            [arg if in_axis is None else _remove_axis(arg, in_axis)
-             for arg, in_axis in zip(args, in_axes)]
-        )
+        args = tuple([
+            arg if in_axis is None else _remove_axis(arg, in_axis)
+            for arg, in_axis in zip(args, in_axes)
+        ])
     stateful_fn.make_jaxpr(state_vals, args)
     return stateful_fn.get_arg_cache_key(state_vals, args)
 
@@ -383,10 +383,7 @@ def _vmap_transform(
         stateful_fn.axis_env = axis_env
 
     # stateful function
-    stateful_fn = StatefulFunction(
-        _vmap_fn_for_compilation,
-        name='vmap',
-    )
+    stateful_fn = StatefulFunction(_vmap_fn_for_compilation, name='vmap')
 
     @functools.wraps(f)
     def new_fn_for_vmap(
@@ -460,7 +457,10 @@ def _vmap_transform(
         # analyze vmapping axis error
         for state in state_trace.get_write_states():
             leaves = jax.tree.leaves(state.value)
-            if any([isinstance(leaf, BatchTracer) for leaf in leaves]) and state not in out_state_to_axis:
+            if (
+                any([isinstance(leaf, BatchTracer) and (leaf.batch_dim is not None) for leaf in leaves])
+                and state not in out_state_to_axis
+            ):
                 if isinstance(state, RandomState) and state in rng_sets:
                     continue
                 state.raise_error_with_source_info(

brainstate/augment/_mapping_test.py

@@ -19,6 +19,8 @@ import unittest
 import jax
 import jax.numpy as jnp
 import numpy as np
+from jax import vmap
+from jax.lax import psum, pmean, pmax
 
 import brainstate
 import brainstate.augment
@@ -433,3 +435,163 @@ class TestVMAPNewStatesEdgeCases(unittest.TestCase):
             foo.c = brainstate.State(jnp.arange(3))  # Original expected shape is (4,)
 
         faulty_init()
+
+
+class TestAxisName:
+    def test1(self):
+        def compute_stats_with_axis_name(x):
+            """Compute statistics using named axis operations"""
+            # Sum across the named axis 'batch'
+            total_sum = psum(x, axis_name='batch')
+
+            # Mean across the named axis 'batch'
+            mean_val = pmean(x, axis_name='batch')
+
+            # Max across the named axis 'batch'
+            max_val = pmax(x, axis_name='batch')
+
+            return {
+                'sum': total_sum,
+                'mean': mean_val,
+                'max': max_val,
+                'original': x
+            }
+
+        batch_data = jnp.array([1.0, 2.0, 3.0, 4.0, 5.0])
+        print("Input batch data:", batch_data)
+
+        # vmap with axis name 'batch'
+        vectorized_stats_jax = jax.jit(vmap(compute_stats_with_axis_name, axis_name='batch'))
+        result_jax = vectorized_stats_jax(batch_data)
+
+        # vmap with axis name 'batch'
+        vectorized_stats = brainstate.transform.vmap(compute_stats_with_axis_name, axis_name='batch')
+        result = vectorized_stats(batch_data)
+
+        # vmap with axis name 'batch'
+        vectorized_stats_v2 = brainstate.transform.jit(
+            brainstate.transform.vmap(compute_stats_with_axis_name, axis_name='batch')
+        )
+        result_v2 = vectorized_stats_v2(batch_data)
+
+        for key in result_jax.keys():
+            print(f"  {key}: {result_jax[key]}")
+            assert jnp.allclose(result_jax[key], result[key]), f"Mismatch in {key}"
+            assert jnp.allclose(result_jax[key], result_v2[key]), f"Mismatch in {key}"
+
+    def test_nested_vmap(self):
+        def nested_computation(x):
+            """Computation with multiple named axes"""
+            # Sum over 'inner' axis, then mean over 'outer' axis
+            inner_sum = psum(x, axis_name='inner')
+            outer_mean = pmean(inner_sum, axis_name='outer')
+            return outer_mean
+
+        # Create 2D batch data
+        data_2d = jnp.arange(12.0).reshape(3, 4)  # Shape: [outer_batch=3, inner_batch=4]
+        print("Input 2D data shape:", data_2d.shape)
+        print("Input 2D data:\n", data_2d)
+
+        # Nested vmap: first over inner dimension, then outer dimension
+        inner_vmap = vmap(nested_computation, axis_name='inner')
+        nested_vmap = vmap(inner_vmap, axis_name='outer')
+
+        result_2d = nested_vmap(data_2d)
+        print("Result after nested vmap:", result_2d)
+
+        inner_vmap_bst = brainstate.transform.vmap(nested_computation, axis_name='inner')
+        nested_vmap_bst = brainstate.transform.vmap(inner_vmap_bst, axis_name='outer')
+        result_2d_bst = nested_vmap_bst(data_2d)
+        print("Result after nested vmap:", result_2d_bst)
+
+        assert jnp.allclose(result_2d, result_2d_bst)
+
+    def _gradient_averaging_simulation_bst(self):
+        def loss_function(params, x, y):
+            """Simple quadratic loss"""
+            pred = params * x
+            return (pred - y) ** 2
+
+        def compute_gradients_with_averaging(params, batch_x, batch_y):
+            """Compute gradients and average them across the batch"""
+            # Compute per-sample gradients
+            grad_fn = jax.grad(loss_function, argnums=0)
+            per_sample_grads = vmap(grad_fn, in_axes=(None, 0, 0))(params, batch_x, batch_y)
+
+            # Average gradients across batch using named axis
+            def average_grads(grads):
+                return pmean(grads, axis_name='batch')
+
+            # Apply averaging with named axis
+            averaged_grads = vmap(average_grads, axis_name='batch')(per_sample_grads)
+            return averaged_grads
+
+        # Example data
+        params = 2.0
+        batch_x = jnp.array([1.0, 2.0, 3.0, 4.0])
+        batch_y = jnp.array([2.0, 4.0, 7.0, 8.0])
+
+        print("Parameters:", params)
+        print("Batch X:", batch_x)
+        print("Batch Y:", batch_y)
+
+        # Compute individual gradients first
+        grad_fn = jax.grad(loss_function, argnums=0)
+        individual_grads = vmap(grad_fn, in_axes=(None, 0, 0))(params, batch_x, batch_y)
+        print("Individual gradients:", individual_grads)
+
+        # Now compute averaged gradients using axis names
+        averaged_grads = compute_gradients_with_averaging(params, batch_x, batch_y)
+        print("Averaged gradients:", averaged_grads)
+
+        return individual_grads, averaged_grads
+
+    def _gradient_averaging_simulation_jax(self):
+        def loss_function(params, x, y):
+            """Simple quadratic loss"""
+            pred = params * x
+            return (pred - y) ** 2
+
+        def compute_gradients_with_averaging(params, batch_x, batch_y):
+            """Compute gradients and average them across the batch"""
+            # Compute per-sample gradients
+            grad_fn = jax.grad(loss_function, argnums=0)
+            per_sample_grads = brainstate.transform.vmap(grad_fn, in_axes=(None, 0, 0))(params, batch_x, batch_y)
+
+            # Average gradients across batch using named axis
+            def average_grads(grads):
+                return pmean(grads, axis_name='batch')
+
+            # Apply averaging with named axis
+            averaged_grads = brainstate.transform.vmap(average_grads, axis_name='batch')(per_sample_grads)
+            return averaged_grads
+
+        # Example data
+        params = 2.0
+        batch_x = jnp.array([1.0, 2.0, 3.0, 4.0])
+        batch_y = jnp.array([2.0, 4.0, 7.0, 8.0])
+
+        print("Parameters:", params)
+        print("Batch X:", batch_x)
+        print("Batch Y:", batch_y)
+
+        # Compute individual gradients first
+        grad_fn = jax.grad(loss_function, argnums=0)
+        individual_grads = brainstate.transform.vmap(grad_fn, in_axes=(None, 0, 0))(params, batch_x, batch_y)
+        print("Individual gradients:", individual_grads)
+
+        # Now compute averaged gradients using axis names
+        averaged_grads = compute_gradients_with_averaging(params, batch_x, batch_y)
+        print("Averaged gradients:", averaged_grads)
+
+        return individual_grads, averaged_grads
+
+    def test_gradient_averaging_simulation(self):
+        individual_grads, averaged_grads = self._gradient_averaging_simulation_bst()
+        individual_grads_jax, averaged_grads_jax = self._gradient_averaging_simulation_jax()
+        assert jnp.allclose(individual_grads, individual_grads_jax)
+        assert jnp.allclose(averaged_grads, averaged_grads_jax)
+
+
+
+

brainstate/nn/_common.py

@@ -118,14 +118,14 @@ class Vmap(Module):
     This class wraps a module and applies vectorized mapping to its execution,
     allowing for efficient parallel processing across specified axes.
 
-    Attributes:
+    Args:
         module (Module): The module to be vmapped.
-        in_axes (int | None | Sequence[Any]): Specifies how to map over inputs.
-        out_axes (Any): Specifies how to map over outputs.
-        vmap_states (Filter | Dict[Filter, int]): Specifies which states to vmap and on which axes.
-        vmap_out_states (Filter | Dict[Filter, int]): Specifies which output states to vmap and on which axes.
-        axis_name (AxisName | None): Name of the axis being mapped over.
-        axis_size (int | None): Size of the axis being mapped over.
+        in_axes (int | None | Sequence[Any], optional): Specifies how to map over inputs. Defaults to 0.
+        out_axes (Any, optional): Specifies how to map over outputs. Defaults to 0.
+        vmap_states (Filter | Dict[Filter, int], optional): Specifies which states to vmap and on which axes. Defaults to None.
+        vmap_out_states (Filter | Dict[Filter, int], optional): Specifies which output states to vmap and on which axes. Defaults to None.
+        axis_name (AxisName | None, optional): Name of the axis being mapped over. Defaults to None.
+        axis_size (int | None, optional): Size of the axis being mapped over. Defaults to None.
     """
 
     def __init__(
@@ -138,18 +138,6 @@ class Vmap(Module):
         axis_name: AxisName | None = None,
         axis_size: int | None = None,
     ):
-        """
-        Initialize the Vmap instance.
-
-        Args:
-            module (Module): The module to be vmapped.
-            in_axes (int | None | Sequence[Any], optional): Specifies how to map over inputs. Defaults to 0.
-            out_axes (Any, optional): Specifies how to map over outputs. Defaults to 0.
-            vmap_states (Filter | Dict[Filter, int], optional): Specifies which states to vmap and on which axes. Defaults to None.
-            vmap_out_states (Filter | Dict[Filter, int], optional): Specifies which output states to vmap and on which axes. Defaults to None.
-            axis_name (AxisName | None, optional): Name of the axis being mapped over. Defaults to None.
-            axis_size (int | None, optional): Size of the axis being mapped over. Defaults to None.
-        """
         super().__init__()
 
         # parameters

brainstate/nn/_exp_euler.py

@@ -69,27 +69,24 @@ def exp_euler_step(
             f'The input data type should be float64, float32, float16, or bfloat16 '
             f'when using Exponential Euler method. But we got {args[0].dtype}.'
         )
+
+    # drift
     dt = environ.get('dt')
     linear, derivative = vector_grad(fn, argnums=0, return_value=True)(*args, **kwargs)
     linear = u.Quantity(u.get_mantissa(linear), u.get_unit(derivative) / u.get_unit(linear))
     phi = u.math.exprel(dt * linear)
     x_next = args[0] + dt * phi * derivative
 
+    # diffusion
     if diffusion is not None:
-        # unit checking
-        diffusion = diffusion(*args, **kwargs)
-        time_unit = u.get_unit(dt)
-        drift_unit = u.get_unit(derivative)
-        diffusion_unit = u.get_unit(diffusion)
-        # if drift_unit.is_unitless:
-        #   assert diffusion_unit.is_unitless, 'The diffusion term should be unitless when the drift term is unitless.'
-        # else:
-        #   u.fail_for_dimension_mismatch(
-        #     drift_unit, diffusion_unit * time_unit ** 0.5,
-        #     "Drift unit is {drift}, diffusion unit is {diffusion}, ",
-        #     drift=drift_unit, diffusion=diffusion_unit * time_unit ** 0.5
-        #   )
-
-        # diffusion
-        x_next += diffusion * u.math.sqrt(dt) * random.randn_like(args[0])
+        diffusion_part = diffusion(*args, **kwargs) * u.math.sqrt(dt) * random.randn_like(args[0])
+        if u.get_dim(x_next) != u.get_dim(diffusion_part):
+            drift_unit = u.get_unit(x_next)
+            time_unit = u.get_unit(dt)
+            raise ValueError(
+                f"Drift unit is {drift_unit}, "
+                f"expected diffusion unit is {drift_unit / time_unit ** 0.5}, "
+                f"but we got {u.get_unit(diffusion_part)}."
+            )
+        x_next += diffusion_part
     return x_next

brainstate/nn/_inputs.py

@@ -547,7 +547,7 @@ def poisson_input(
                 num_input,
                 p,
                 tar[indices].shape,
-                # check_valid=False,
+                check_valid=False,
                 dtype=tar.dtype
             ),
             tar_val,

brainstate/random/_rand_state.py

@@ -384,7 +384,10 @@ class RandomState(State):
         loc = _check_py_seq(loc)
         scale = _check_py_seq(scale)
         if size is None:
-            size = lax.broadcast_shapes(jnp.shape(loc), jnp.shape(scale))
+            size = lax.broadcast_shapes(
+                jnp.shape(loc) if loc is not None else (),
+                jnp.shape(scale) if scale is not None else ()
+            )
         key = self.split_key() if key is None else _formalize_key(key)
         dtype = dtype or environ.dftype()
         r = _loc_scale(loc, scale, jr.logistic(key, shape=_size2shape(size), dtype=dtype))
@@ -399,7 +402,10 @@ class RandomState(State):
         loc = _check_py_seq(loc)
         scale = _check_py_seq(scale)
         if size is None:
-            size = lax.broadcast_shapes(jnp.shape(scale), jnp.shape(loc))
+            size = lax.broadcast_shapes(
+                jnp.shape(scale) if scale is not None else (),
+                jnp.shape(loc) if loc is not None else ()
+            )
         key = self.split_key() if key is None else _formalize_key(key)
         dtype = dtype or environ.dftype()
         r = _loc_scale(loc, scale, jr.normal(key, shape=_size2shape(size), dtype=dtype))
@@ -456,7 +462,7 @@ class RandomState(State):
                        dtype: DTypeLike = None):
         shape = _check_py_seq(shape)
         if size is None:
-            size = jnp.shape(shape)
+            size = jnp.shape(shape) if shape is not None else ()
         key = self.split_key() if key is None else _formalize_key(key)
         dtype = dtype or environ.dftype()
         r = jr.gamma(key, a=shape, shape=_size2shape(size), dtype=dtype)
@@ -477,7 +483,7 @@ class RandomState(State):
                    dtype: DTypeLike = None):
         df = _check_py_seq(df)
         if size is None:
-            size = jnp.shape(size)
+            size = jnp.shape(size) if size is not None else ()
         key = self.split_key() if key is None else _formalize_key(key)
         dtype = dtype or environ.dftype()
         r = jr.t(key, df=df, shape=_size2shape(size), dtype=dtype)
@@ -606,8 +612,8 @@ class RandomState(State):
 
         if size is None:
             size = jnp.broadcast_shapes(
-                jnp.shape(mean),
-                jnp.shape(sigma)
+                jnp.shape(mean) if mean is not None else (),
+                jnp.shape(sigma) if sigma is not None else ()
             )
         key = self.split_key() if key is None else _formalize_key(key)
         dtype = dtype or environ.dftype()
@@ -822,7 +828,7 @@ class RandomState(State):
         a = _check_py_seq(a)
         scale = _check_py_seq(scale)
         if size is None:
-            size = jnp.broadcast_shapes(jnp.shape(a), jnp.shape(scale))
+            size = jnp.broadcast_shapes(jnp.shape(a), jnp.shape(scale) if scale is not None else ())
         else:
             if jnp.size(a) > 1:
                 raise ValueError(f'"a" should be a scalar when "size" is provided. But we got {a}')

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.1.4
+Version: 0.1.5
 Summary: A ``State``-based Transformation System for Program Compilation and Augmentation.
 Home-page: https://github.com/chaobrain/brainstate
 Author: BrainState Developers

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

@@ -1,6 +1,6 @@
-brainstate/__init__.py,sha256=A_OD4cJiVu3xpthNGJh6fhRjCKsI7_Mxsow3Al4m2-w,1496
+brainstate/__init__.py,sha256=24AzLahGzsgxErIIbDuFSzGFn-btHTh9dyczmpOkOvA,1496
 brainstate/_compatible_import.py,sha256=LUSZlA0APWozxM8Kf9pZrM2YbwY7X3jzVVHZInaBL7Y,4630
-brainstate/_state.py,sha256=o5Kk4bGwVz6Dfj9dlmZqdh6zUXcz6Tvc6WOjH9ajlIU,60716
+brainstate/_state.py,sha256=OBQIu1eLH4szHw2SGduPxx24gKoU7jXgCL-TPNdI_iw,60735
 brainstate/_state_test.py,sha256=b6uvZdVRyC4n6-fYzmHNry1b-gJ6zE_kRSxGinqiHaw,1638
 brainstate/_utils.py,sha256=j-b239RHfC5BnvhGbSExQpdY21LrMEyWMSHBdNGThOI,1657
 brainstate/environ.py,sha256=VgtG0S_aR1g_1gplRWg_v2ZrcS-F6LZk35BPCBgsIIA,17660
@@ -15,8 +15,8 @@ brainstate/augment/_autograd.py,sha256=zMMLZOidQq2p96wzEOgR-MynV2JH1l1AnoJ28eVwD
 brainstate/augment/_autograd_test.py,sha256=UhNd41luca_Kj9a8byL3Dq64Ta55WU-je-LCiK_z0Vc,45060
 brainstate/augment/_eval_shape.py,sha256=mYRB5clSRHdXX0c4c5_WDIIBtRl4a9-xl8Tg60mgLBI,3854
 brainstate/augment/_eval_shape_test.py,sha256=WXrmZKmnykmYPveRfZbrDF0sFm2_PTr982yl4JM7Ebg,1390
-brainstate/augment/_mapping.py,sha256=AvGt8qhOdBgRJyWZn7J5JnOvxZc7XfZ_JKXaECOQjIw,43500
-brainstate/augment/_mapping_test.py,sha256=An5iSqrAXq_OVTvY3GFTjsB9AvFZ4j6AWNZQ_WWfra4,15319
+brainstate/augment/_mapping.py,sha256=M-xl7Q-wGrftI0fJ4jO-SzCEpB7XZ_UjQfkJs_2TbCQ,43557
+brainstate/augment/_mapping_test.py,sha256=tEyXioA4_v4WGRFrpgUInyKyBjAVNLonWdHKVacmphU,21946
 brainstate/augment/_random.py,sha256=bkngsIk6Wwi3e46I7YSbBjLCGAz0Q3WuadUH4mqTjbY,5348
 brainstate/compile/__init__.py,sha256=fQtG316MLkeeu1Ssp54Kghw1PwbGK5gNq9yRVJu0wjA,1474
 brainstate/compile/_ad_checkpoint.py,sha256=HM6L90HU0N84S30uJpX8wqTO0ZcDnctqmwf8qFlkDYo,9358
@@ -57,7 +57,7 @@ brainstate/init/_regular_inits_test.py,sha256=bvv0AOfLOEP0BIQIBLztKw3EPyEp7n2fHW
 brainstate/nn/__init__.py,sha256=nDIBDlzHaBy-vk4cb7FiPO6zoNe2tI5qvHIK8O7yrOU,3721
 brainstate/nn/_collective_ops.py,sha256=rWcjqaP0rW6sXdhE0fjtDi4twmq5VvhNEg8AoiQ-tDU,17390
 brainstate/nn/_collective_ops_test.py,sha256=bwq0DApcsk0_2xpxMl0_e2cGKT63g5rSngpigCm07ps,1409
-brainstate/nn/_common.py,sha256=Pt8P1LgE0qW3QnfX8CQQGH3yZ78RL7NkIqqYEw_z8xs,7096
+brainstate/nn/_common.py,sha256=qHAOID_eeKiPUXk_ION65sbYXyF-ddH5w5BayvH8Thg,6431
 brainstate/nn/_conv.py,sha256=Zk-yj34n6CkjntcM9xpMGLTxKNfWdIWsTsoGbtdL0yU,18448
 brainstate/nn/_conv_test.py,sha256=2lcUTG7twkyhuyKwuBux-NgU8NU_W4Cp1-G8EyDJ_uk,8862
 brainstate/nn/_delay.py,sha256=l36FBgNhfL64tM3VGOsJNTtKr44HjxxtBWMFFCm3Pks,17361
@@ -68,11 +68,11 @@ brainstate/nn/_dynamics_test.py,sha256=w7AV57LdhbBNYprdFpKq8MFSCbXKVkGgp_NbL3ANX
 brainstate/nn/_elementwise.py,sha256=4czeJWGQopV49iZo8DuN_WzAbXoMC1gtqaGjlON6e7c,33291
 brainstate/nn/_elementwise_test.py,sha256=_dd9eX2ZJ7p24ahuoapCaRTZ0g1boufXMyqHFx1d4WY,5688
 brainstate/nn/_embedding.py,sha256=SaAJbgXmuJ8XlCOX9ob4yvmgh9Fk627wMguRzJMJ1H8,2138
-brainstate/nn/_exp_euler.py,sha256=ndDB43PM4jsZKu_zdLTZ2-ojnuNrg55LZap23oBTtdA,3493
+brainstate/nn/_exp_euler.py,sha256=WTpZm-XQmsdMLNazY7wIu8eeO6pK0kRzt2lJnhEgMIk,3293
 brainstate/nn/_exp_euler_test.py,sha256=XD--qMbGHrHa3WtcPMmJKk59giDcEhSqZuBOmTNYUr8,1227
 brainstate/nn/_fixedprob.py,sha256=KGXohiU0wZnFIQDuwiRUTFsbsr8R0p8zgi5UZDuv1Bk,10004
 brainstate/nn/_fixedprob_test.py,sha256=qbRBh-MpMtEOsg492gFu2w9-FOP9z_bXapm-Q0gLLYM,3929
-brainstate/nn/_inputs.py,sha256=wPOfPE4IesNoDmxZJxqR0siBlJioEX-_1IZ2cltAIpM,20605
+brainstate/nn/_inputs.py,sha256=hMPkx9qDBpJWPshZXLF4H1QiYK1-46wntHUIlG7cT7c,20603
 brainstate/nn/_linear.py,sha256=5WuhcqU-uBUC91vnwezQYMHPKmlZPDgIJ5UpffxoX1I,14472
 brainstate/nn/_linear_mv.py,sha256=6hDXx4yPqRSa7uIsW9f9eJuy23dcXN9Mp2_lSvw8BDA,2635
 brainstate/nn/_linear_mv_test.py,sha256=ZCM1Zy6mImQfCfdZOGnTwkiLLPXK5yalv1Ts9sWZuPA,3864
@@ -111,7 +111,7 @@ brainstate/random/_rand_funs.py,sha256=c4xiY2NeMizSslxbWOa-QJZ3h-LDfsgH4fbvBMSNL
 brainstate/random/_rand_funs_test.py,sha256=G8BuxDjBSeE-Mh7KuwVk6mPQhHE0m1R5HDFljEOIdzg,20669
 brainstate/random/_rand_seed.py,sha256=1ZdfFZWyOhpd72EDdEDmpkp3yoLVwdv-sGI9BwiZfzI,5949
 brainstate/random/_rand_seed_test.py,sha256=waXXfch57X1XE1zDnCRokT6ziZOK0g-lYE80o6epDYM,1536
-brainstate/random/_rand_state.py,sha256=dTO7wqmuTYRdPy7ItsrK-7aNt5QQXOsZ4XwiH7mzmy8,55170
+brainstate/random/_rand_state.py,sha256=hQ_govRxfIgsmtFe_V2B-jftqiMWRWrd9wkviPFOziY,55523
 brainstate/random/_random_for_unit.py,sha256=kGp4EUX19MXJ9Govoivbg8N0bddqOldKEI2h_TbdONY,2057
 brainstate/util/__init__.py,sha256=6efwr63osmqviNU_6_Nufag19PwxRDFvDAZrq6sH5yo,1555
 brainstate/util/_pretty_pytree_test.py,sha256=Dn0TdjX6wLBXaTD4jfYTu6cKfFHwKSxi4_3bX7kB_IA,5621
@@ -124,8 +124,8 @@ brainstate/util/pretty_repr.py,sha256=7Xp7IFNUeP7cGlpvwwJyBslbQVnXEqC1I6neV1Jx1S
 brainstate/util/pretty_table.py,sha256=uJVaamFGQ4nKP8TkEGPWXHpzjMecDo2q1Ah6XtRjdPY,108117
 brainstate/util/scaling.py,sha256=U6DM-afPrLejiGqo1Nla7z4YbTBVicctsBEweurr_mk,7524
 brainstate/util/struct.py,sha256=7HbbQNrZ3zxYw93MU1bUZ9ZPBKftYOVKuXEochLSErw,17479
-brainstate-0.1.4.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.1.4.dist-info/METADATA,sha256=EWarKpuYFIN0ITFn8lwsUy0h9-KjrVPDwcvjQMHFRto,4135
-brainstate-0.1.4.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.1.4.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.1.4.dist-info/RECORD,,
+brainstate-0.1.5.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.1.5.dist-info/METADATA,sha256=0rd8rcVbLHdAkaLX9T2-BA2aRfq4D2ToDcmt0pp0wwg,4135
+brainstate-0.1.5.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+brainstate-0.1.5.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.1.5.dist-info/RECORD,,