brainstate 0.1.0.post20241129__py2.py3-none-any.whl → 0.1.0.post20241210__py2.py3-none-any.whl

brainstate/_state.py

@@ -99,6 +99,9 @@ def catch_new_states(tag: str = None) -> List:
 
 
 class Catcher:
+    """
+    The catcher to catch the new states.
+    """
     def __init__(self, tag: str):
         self.tag = tag
         self.state_ids = set()
@@ -231,6 +234,7 @@ class State(Generic[A], PrettyRepr):
         # avoid using self._setattr to avoid the check
         vars(self).update(metadata)
 
+        # record the state initialization
         record_state_init(self)
 
     if not TYPE_CHECKING:
@@ -290,7 +294,6 @@ class State(Generic[A], PrettyRepr):
           v: The value.
         """
         self.write_value(v)
-        self._been_writen = True
 
     def write_value(self, v) -> None:
         # value checking
@@ -301,6 +304,8 @@ class State(Generic[A], PrettyRepr):
         record_state_value_write(self)
         # set the value
         self._value = v
+        # set flag
+        self._been_writen = True
 
     def restore_value(self, v) -> None:
         """
@@ -511,6 +516,15 @@ class LongTermState(State):
     __module__ = 'brainstate'
 
 
+class BatchState(LongTermState):
+    """
+    The batch state, which is used to store the batch data in the program.
+    """
+
+    __module__ = 'brainstate'
+
+
+
 class HiddenState(ShortTermState):
     """
     The hidden state, which is used to store the hidden data in a dynamic model.

brainstate/augment/_mapping_test.py

@@ -204,7 +204,91 @@ class TestVmap(unittest.TestCase):
         self.assertTrue(jnp.allclose(foo.a.value, jnp.arange(4)))
         self.assertTrue(jnp.allclose(foo.b.value, jnp.arange(4) * jnp.arange(4)))
 
+        write_state_ids = [id(st) for st in trace.get_write_states()]
+        read_state_ids = [id(st) for st in trace.get_read_states()]
+
+        assert id(foo.a) in read_state_ids
+        assert id(foo.b) in write_state_ids
+
+        print(trace.get_write_states())
+        print(trace.get_read_states())
+
+
+
+    def test_vmap_jit(self):
+        class Foo(bst.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.a = bst.ParamState(jnp.arange(4))
+                self.b = bst.ShortTermState(jnp.arange(4))
+
+            def __call__(self):
+                self.b.value = self.a.value * self.b.value
+
+        @bst.augment.vmap
+        def mul(foo):
+            foo()
+
+        @bst.compile.jit
+        def mul_jit(inp):
+            mul(foo)
+            foo.a.value += inp
+
+        foo = Foo()
+        with bst.StateTraceStack() as trace:
+            mul_jit(1.)
+
+        print(foo.a.value)
+        print(foo.b.value)
+        self.assertTrue(jnp.allclose(foo.a.value, jnp.arange(4) + 1.))
+        self.assertTrue(jnp.allclose(foo.b.value, jnp.arange(4) * jnp.arange(4)))
+
+        write_state_ids = [id(st) for st in trace.get_write_states()]
+        read_state_ids = [id(st) for st in trace.get_read_states()]
+
+        assert id(foo.a) in write_state_ids
+        assert id(foo.b) in write_state_ids
+
         print(trace.get_write_states())
         print(trace.get_read_states())
 
 
+    def test_vmap_grad(self):
+        class Foo(bst.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.a = bst.ParamState(jnp.arange(4.))
+                self.b = bst.ShortTermState(jnp.arange(4.))
+
+            def __call__(self):
+                self.b.value = self.a.value * self.b.value
+
+        @bst.augment.vmap
+        def mul(foo):
+            foo()
+
+        def loss():
+            mul(foo)
+            return jnp.sum(foo.b.value)
+
+        foo = Foo()
+        with bst.StateTraceStack() as trace:
+            grads, loss = bst.augment.grad(loss, foo.states(bst.ParamState), return_value=True)()
+            print(grads)
+            print(loss)
+
+        # print(foo.a.value)
+        # print(foo.b.value)
+        # self.assertTrue(jnp.allclose(foo.a.value, jnp.arange(4) + 1.))
+        # self.assertTrue(jnp.allclose(foo.b.value, jnp.arange(4) * jnp.arange(4)))
+        #
+        # write_state_ids = [id(st) for st in trace.get_write_states()]
+        # read_state_ids = [id(st) for st in trace.get_read_states()]
+        #
+        # assert id(foo.a) in write_state_ids
+        # assert id(foo.b) in write_state_ids
+        #
+        # print(trace.get_write_states())
+        # print(trace.get_read_states())
+
+

brainstate/compile/_conditions.py

@@ -94,9 +94,8 @@ def cond(pred, true_fun: Callable, false_fun: Callable, *operands):
             return false_fun(*operands)
 
     # evaluate jaxpr
-    with jax.ensure_compile_time_eval():
-        stateful_true = StatefulFunction(true_fun).make_jaxpr(*operands)
-        stateful_false = StatefulFunction(false_fun).make_jaxpr(*operands)
+    stateful_true = StatefulFunction(true_fun).make_jaxpr(*operands)
+    stateful_false = StatefulFunction(false_fun).make_jaxpr(*operands)
 
     # state trace and state values
     state_trace = stateful_true.get_state_trace() + stateful_false.get_state_trace()
@@ -175,10 +174,9 @@ def switch(index, branches: Sequence[Callable], *operands):
         return branches[int(index)](*operands)
 
     # evaluate jaxpr
-    with jax.ensure_compile_time_eval():
-        wrapped_branches = [StatefulFunction(branch) for branch in branches]
-        for wrapped_branch in wrapped_branches:
-            wrapped_branch.make_jaxpr(*operands)
+    wrapped_branches = [StatefulFunction(branch) for branch in branches]
+    for wrapped_branch in wrapped_branches:
+        wrapped_branch.make_jaxpr(*operands)
 
     # wrap the functions
     state_trace = wrapped_branches[0].get_state_trace() + wrapped_branches[1].get_state_trace()

brainstate/compile/_jit.py

@@ -83,9 +83,9 @@ def _get_jitted_fun(
             return fun.fun(*args, **params)
 
         # compile the function and get the state trace
-        with jax.ensure_compile_time_eval():
-            state_trace = fun.compile_function_and_get_state_trace(*args, **params, return_only_write=True)
-            read_state_vals = state_trace.get_read_state_values(True)
+        state_trace = fun.compile_function_and_get_state_trace(*args, **params, return_only_write=True)
+        read_state_vals = state_trace.get_read_state_values(True)
+
         # call the jitted function
         write_state_vals, outs = jit_fun(state_trace.get_state_values(), *args, **params)
         # write the state values back to the states

brainstate/compile/_loop_collect_return.py

@@ -202,16 +202,19 @@ def scan(
     # ------------------------------ #
     xs_avals = [jax.core.raise_to_shaped(jax.core.get_aval(x)) for x in xs_flat]
     x_avals = [jax.core.mapped_aval(length, 0, aval) for aval in xs_avals]
-    with jax.ensure_compile_time_eval():
-        stateful_fun = StatefulFunction(f).make_jaxpr(init, xs_tree.unflatten(x_avals))
-        state_trace = stateful_fun.get_state_trace()
-        all_writen_state_vals = state_trace.get_write_state_values(True)
-        all_read_state_vals = state_trace.get_read_state_values(True)
-        wrapped_f = wrap_single_fun(stateful_fun, state_trace.been_writen, all_read_state_vals)
+    stateful_fun = StatefulFunction(f).make_jaxpr(init, xs_tree.unflatten(x_avals))
+    state_trace = stateful_fun.get_state_trace()
+    all_writen_state_vals = state_trace.get_write_state_values(True)
+    all_read_state_vals = state_trace.get_read_state_values(True)
+    wrapped_f = wrap_single_fun(stateful_fun, state_trace.been_writen, all_read_state_vals)
 
     # scan
     init = (all_writen_state_vals, init)
-    (all_writen_state_vals, carry), ys = jax.lax.scan(wrapped_f, init, xs, length=length, reverse=reverse,
+    (all_writen_state_vals, carry), ys = jax.lax.scan(wrapped_f,
+                                                      init,
+                                                      xs,
+                                                      length=length,
+                                                      reverse=reverse,
                                                       unroll=unroll)
     # assign the written state values and restore the read state values
     write_back_state_values(state_trace, all_read_state_vals, all_writen_state_vals)

brainstate/compile/_loop_no_collection.py

@@ -103,11 +103,10 @@ def while_loop(
             pass
 
     # evaluate jaxpr
-    with jax.ensure_compile_time_eval():
-        stateful_cond = StatefulFunction(cond_fun).make_jaxpr(init_val)
-        stateful_body = StatefulFunction(body_fun).make_jaxpr(init_val)
-        if len(stateful_cond.get_write_states()) != 0:
-            raise ValueError("while_loop: cond_fun should not have any write states.")
+    stateful_cond = StatefulFunction(cond_fun).make_jaxpr(init_val)
+    stateful_body = StatefulFunction(body_fun).make_jaxpr(init_val)
+    if len(stateful_cond.get_write_states()) != 0:
+        raise ValueError("while_loop: cond_fun should not have any write states.")
 
     # state trace and state values
     state_trace = stateful_cond.get_state_trace() + stateful_body.get_state_trace()

brainstate/compile/_make_jaxpr.py

@@ -72,7 +72,6 @@ from jax.util import wraps
 from brainstate._state import State, StateTraceStack
 from brainstate._utils import set_module_as
 from brainstate.typing import PyTree
-from brainstate.util._tracers import new_jax_trace
 
 AxisName = Hashable
 
@@ -112,28 +111,27 @@ def _new_arg_fn(frame, trace, aval):
     return tracer
 
 
-def _init_state_trace() -> StateTraceStack:
-    # Should be within the calling of ``jax.make_jaxpr()``
-    frame, trace = new_jax_trace()
+def _new_jax_trace():
+    main = jax.core.thread_local_state.trace_state.trace_stack.stack[-1]
+    frame = main.jaxpr_stack[-1]
+    trace = pe.DynamicJaxprTrace(main, jax.core.cur_sublevel())
+    return frame, trace
+
+
+def _init_state_trace_stack() -> StateTraceStack:
     state_trace: StateTraceStack = StateTraceStack()
-    # Set the function to transform the new argument to a tracer
-    state_trace.set_new_arg(functools.partial(_new_arg_fn, frame, trace))
-    return state_trace
 
+    if jax.__version_info__ < (0, 4, 36):
+        # Should be within the calling of ``jax.make_jaxpr()``
+        frame, trace = _new_jax_trace()
+        # Set the function to transform the new argument to a tracer
+        state_trace.set_new_arg(functools.partial(_new_arg_fn, frame, trace))
+        return state_trace
 
-# def wrapped_abstractify(x: Any) -> Any:
-#   """
-#   Abstractify the input.
-#
-#   Args:
-#     x: The input.
-#
-#   Returns:
-#     The abstractified input.
-#   """
-#   if isinstance(x, pe.DynamicJaxprTracer):
-#     return jax.core.ShapedArray(x.aval.shape, x.aval.dtype, weak_type=x.aval.weak_type)
-#   return shaped_abstractify(x)
+    else:
+        trace = jax.core.trace_ctx.trace
+        state_trace.set_new_arg(trace.new_arg)
+        return state_trace
 
 
 class StatefulFunction(object):
@@ -383,12 +381,15 @@ class StatefulFunction(object):
           A tuple of the states that are read and written by the function and the output of the function.
         """
         # state trace
-        state_trace = _init_state_trace()
+        state_trace = _init_state_trace_stack()
         self._cached_state_trace[cache_key] = state_trace
         with state_trace:
             out = self.fun(*args, **kwargs)
-            state_values = state_trace.get_write_state_values(
-                True) if return_only_write else state_trace.get_state_values()
+            state_values = (
+                state_trace.get_write_state_values(True)
+                if return_only_write else
+                state_trace.get_state_values()
+            )
         state_trace.recovery_original_values()
 
         # State instance as functional returns is not allowed.
@@ -419,17 +420,21 @@ class StatefulFunction(object):
             try:
                 # jaxpr
                 jaxpr, (out_shapes, state_shapes) = _make_jaxpr(
-                    functools.partial(self._wrapped_fun_to_eval, cache_key, return_only_write=return_only_write),
+                    functools.partial(
+                        self._wrapped_fun_to_eval,
+                        cache_key,
+                        return_only_write=return_only_write
+                    ),
                     static_argnums=self.static_argnums,
                     axis_env=self.axis_env,
                     return_shape=True,
                     abstracted_axes=self.abstracted_axes
                 )(*args, **kwargs)
-
                 # returns
                 self._cached_jaxpr_out_tree[cache_key] = jax.tree.structure((out_shapes, state_shapes))
                 self._cached_out_shapes[cache_key] = (out_shapes, state_shapes)
                 self._cached_jaxpr[cache_key] = jaxpr
+
             except Exception as e:
                 try:
                     self._cached_state_trace.pop(cache_key)

brainstate/compile/_progress_bar.py

@@ -93,14 +93,32 @@ class ProgressBarRunner(object):
             self.tqdm_bars[0].update(self.remainder)
         self.tqdm_bars[0].close()
 
+    def _tqdm(self, is_init, is_print, is_final):
+        if is_init:
+            self.tqdm_bars[0] = tqdm(range(self.n), **self.kwargs)
+            self.tqdm_bars[0].set_description(self.message, refresh=False)
+        if is_print:
+            self.tqdm_bars[0].update(self.print_freq)
+        if is_final:
+            if self.remainder > 0:
+                self.tqdm_bars[0].update(self.remainder)
+            self.tqdm_bars[0].close()
+
     def __call__(self, iter_num, *args, **kwargs):
+        # jax.debug.callback(
+        #     self._tqdm,
+        #     iter_num == 0,
+        #     (iter_num + 1) % self.print_freq == 0,
+        #     iter_num == self.n - 1
+        # )
+
         _ = jax.lax.cond(
             iter_num == 0,
             lambda: jax.debug.callback(self._define_tqdm),
             lambda: None,
         )
         _ = jax.lax.cond(
-            (iter_num + 1) % self.print_freq == 0,
+            iter_num % self.print_freq == (self.print_freq - 1),
             lambda: jax.debug.callback(self._update_tqdm),
             lambda: None,
         )
@@ -109,3 +127,4 @@ class ProgressBarRunner(object):
             lambda: jax.debug.callback(self._close_tqdm),
             lambda: None,
         )
+

brainstate/functional/_activations.py

@@ -588,8 +588,7 @@ def glu(x: ArrayLike, axis: int = -1) -> Union[jax.Array, u.Quantity]:
 
 def log_softmax(x: ArrayLike,
                 axis: int | tuple[int, ...] | None = -1,
-                where: ArrayLike | None = None,
-                initial: ArrayLike | None = None) -> Union[jax.Array, u.Quantity]:
+                where: ArrayLike | None = None) -> Union[jax.Array, u.Quantity]:
     r"""Log-Softmax function.
 
     Computes the logarithm of the :code:`softmax` function, which rescales
@@ -604,8 +603,6 @@ def log_softmax(x: ArrayLike,
       axis: the axis or axes along which the :code:`log_softmax` should be
         computed. Either an integer or a tuple of integers.
       where: Elements to include in the :code:`log_softmax`.
-      initial: The minimum value used to shift the input array. Must be present
-        when :code:`where` is not None.
 
     Returns:
       An array.
@@ -613,15 +610,12 @@ def log_softmax(x: ArrayLike,
     See also:
       :func:`softmax`
     """
-    if initial is not None:
-        initial = u.Quantity(initial).in_unit(u.get_unit(x)).mantissa
-    return _keep_unit(jax.nn.log_softmax, x, axis=axis, where=where, initial=initial)
+    return _keep_unit(jax.nn.log_softmax, x, axis=axis, where=where)
 
 
 def softmax(x: ArrayLike,
             axis: int | tuple[int, ...] | None = -1,
-            where: ArrayLike | None = None,
-            initial: ArrayLike | None = None) -> Union[jax.Array, u.Quantity]:
+            where: ArrayLike | None = None) -> Union[jax.Array, u.Quantity]:
     r"""Softmax function.
 
     Computes the function which rescales elements to the range :math:`[0, 1]`
@@ -645,9 +639,7 @@ def softmax(x: ArrayLike,
     See also:
       :func:`log_softmax`
     """
-    if initial is not None:
-        initial = u.Quantity(initial).in_unit(u.get_unit(x)).mantissa
-    return _keep_unit(jax.nn.softmax, x, axis=axis, where=where, initial=initial)
+    return _keep_unit(jax.nn.softmax, x, axis=axis, where=where)
 
 
 def standardize(x: ArrayLike,

brainstate/graph/_graph_operation.py

@@ -608,7 +608,10 @@ def _get_children(graph_def, state_mapping, index_ref, index_ref_cache):
                         if isinstance(value, TreefyState):
                             variable.update_from_ref(value)
                         elif isinstance(value, State):
-                            variable.restore_value(value.value)
+                             if value._been_writen:
+                                variable.write_value(value.value)
+                             else:
+                                variable.restore_value(value.value)
                         else:
                             raise ValueError(f'Expected a State type for {key!r}, but got {type(value)}.')
                     else:  # if it doesn't, create a new variable

brainstate/nn/_collective_ops.py

@@ -20,8 +20,10 @@ from typing import Dict, Callable, TypeVar
 
 import jax
 
+from brainstate._state import catch_new_states
 from brainstate._utils import set_module_as
 from brainstate.graph import nodes
+from brainstate.util._filter import Filter
 from ._module import Module
 
 # the maximum order
@@ -74,16 +76,29 @@ def call_order(level: int = 0, check_order_boundary: bool = True):
 
 
 @set_module_as('brainstate.nn')
-def init_all_states(target: T, *args, exclude=None, **kwargs) -> T:
+def init_all_states(
+    target: T,
+    *args,
+    exclude: Filter = None,
+    **kwargs
+) -> T:
     """
     Collectively initialize states of all children nodes in the given target.
 
     Args:
       target: The target Module.
+      exclude: The filter to exclude some nodes.
+      tag: The tag for the new states.
+      args: The positional arguments for the initialization, which will be passed to the `init_state` method
+        of each node.
+      kwargs: The keyword arguments for the initialization, which will be passed to the `init_state` method
+        of each node.
 
     Returns:
       The target Module.
     """
+
+    # node that has `call_order` decorated
     nodes_with_order = []
 
     nodes_ = nodes(target).filter(Module)
@@ -97,7 +112,7 @@ def init_all_states(target: T, *args, exclude=None, **kwargs) -> T:
         else:
             node.init_state(*args, **kwargs)
 
-    # reset the node's states
+    # reset the node's states with `call_order`
     for node in sorted(nodes_with_order, key=lambda x: x.init_state.call_order):
         node.init_state(*args, **kwargs)
 
@@ -115,6 +130,7 @@ def reset_all_states(target: Module, *args, **kwargs) -> Module:
     Returns:
       The target Module.
     """
+
     nodes_with_order = []
 
     # reset node whose `init_state` has no `call_order`

brainstate/nn/_dyn_impl/_dynamics_synapse.py

@@ -112,7 +112,7 @@ class STP(Synapse):
         self.u.value = init.param(init.Constant(self.U), self.varshape, batch_size)
 
     def update(self, pre_spike):
-        u = exp_euler_step(lambda u: self.U - u / self.tau_f, self.u.value)
+        u = exp_euler_step(lambda u: - u / self.tau_f, self.u.value)
         x = exp_euler_step(lambda x: (1 - x) / self.tau_d, self.x.value)
 
         # --- original code:

brainstate/nn/_elementwise/_dropout.py

@@ -17,7 +17,7 @@
 from __future__ import annotations
 
 from functools import partial
-from typing import Optional
+from typing import Optional, Sequence
 
 import brainunit as u
 import jax.numpy as jnp
@@ -29,7 +29,6 @@ from brainstate.typing import Size
 
 __all__ = [
     'DropoutFixed', 'Dropout', 'Dropout1d', 'Dropout2d', 'Dropout3d',
-    'AlphaDropout', 'FeatureAlphaDropout',
 ]
 
 
@@ -47,9 +46,9 @@ class Dropout(ElementWiseBlock):
            research 15.1 (2014): 1929-1958.
 
     Args:
-      prob: Probability to keep element of the tensor.
-      mode: Mode. The computation mode of the object.
-      name: str. The name of the dynamic system.
+        prob: Probability to keep element of the tensor.
+        broadcast_dims: dimensions that will share the same dropout mask.
+        name: str. The name of the dynamic system.
 
     """
     __module__ = 'brainstate.nn'
@@ -57,20 +56,28 @@ class Dropout(ElementWiseBlock):
     def __init__(
         self,
         prob: float = 0.5,
+        broadcast_dims: Sequence[int] = (),
         name: Optional[str] = None
     ) -> None:
         super().__init__(name=name)
         assert 0. <= prob <= 1., f"Dropout probability must be in the range [0, 1]. But got {prob}."
         self.prob = prob
+        self.broadcast_dims = broadcast_dims
 
     def __call__(self, x):
         dtype = u.math.get_dtype(x)
         fit_phase = environ.get('fit', desc='Whether this is a fitting process. Bool.')
         if fit_phase and self.prob < 1.:
-            keep_mask = random.bernoulli(self.prob, x.shape)
-            return jnp.where(keep_mask,
-                             jnp.asarray(x / self.prob, dtype=dtype),
-                             jnp.asarray(0., dtype=dtype))
+            broadcast_shape = list(x.shape)
+            for dim in self.broadcast_dims:
+                broadcast_shape[dim] = 1
+            keep_mask = random.bernoulli(self.prob, broadcast_shape)
+            keep_mask = jnp.broadcast_to(keep_mask, x.shape)
+            return jnp.where(
+                keep_mask,
+                jnp.asarray(x / self.prob, dtype=dtype),
+                jnp.asarray(0., dtype=dtype)
+            )
         else:
             return x
 
@@ -93,7 +100,6 @@ class _DropoutNd(ElementWiseBlock):
         self.channel_axis = channel_axis
 
     def __call__(self, x):
-
         # check input shape
         inp_dim = u.math.ndim(x)
         if inp_dim not in (self.minimal_dim, self.minimal_dim + 1):
@@ -114,10 +120,13 @@ class _DropoutNd(ElementWiseBlock):
         # generate mask
         if fit_phase and self.prob < 1.:
             dtype = u.math.get_dtype(x)
-            keep_mask = jnp.broadcast_to(random.bernoulli(self.prob, mask_shape), x.shape)
-            return jnp.where(keep_mask,
-                             jnp.asarray(x / self.prob, dtype=dtype),
-                             jnp.asarray(0., dtype=dtype))
+            keep_mask = random.bernoulli(self.prob, mask_shape)
+            keep_mask = jnp.broadcast_to(keep_mask, x.shape)
+            return jnp.where(
+                keep_mask,
+                jnp.asarray(x / self.prob, dtype=dtype),
+                jnp.asarray(0., dtype=dtype)
+            )
         else:
             return x
 
@@ -296,8 +305,8 @@ class AlphaDropout(_DropoutNd):
     """
     __module__ = 'brainstate.nn'
 
-    def forward(self, x):
-        return F.alpha_dropout(x, self.p, self.training)
+    def update(self, *args, **kwargs):
+        raise NotImplementedError("AlphaDropout is not supported in the current version.")
 
 
 class FeatureAlphaDropout(_DropoutNd):
@@ -344,8 +353,8 @@ class FeatureAlphaDropout(_DropoutNd):
     """
     __module__ = 'brainstate.nn'
 
-    def forward(self, x):
-        return F.feature_alpha_dropout(x, self.p, self.training)
+    def update(self, *args, **kwargs):
+        raise NotImplementedError("FeatureAlphaDropout is not supported in the current version.")
 
 
 class DropoutFixed(ElementWiseBlock):

brainstate/nn/_interaction/_normalizations.py

@@ -23,7 +23,7 @@ import jax
 import jax.numpy as jnp
 
 from brainstate import environ, init
-from brainstate._state import LongTermState, ParamState
+from brainstate._state import ParamState, BatchState
 from brainstate.nn._module import Module
 from brainstate.typing import DTypeLike, ArrayLike, Size, Axes
 
@@ -91,6 +91,18 @@ def _abs_sq(x):
         return jax.lax.square(x)
 
 
+class NormalizationParamState(ParamState):
+    # This is a dummy class to be used as a compatibility
+    # usage of `ETraceParam` for the layers in "brainetrace"
+    def execute(self, x):
+        param = self.value
+        if 'scale' in param:
+            x = x * param['scale']
+        if 'bias' in param:
+            x = x + param['bias']
+        return x
+
+
 def _compute_stats(
     x: ArrayLike,
     axes: Sequence[int],
@@ -150,12 +162,17 @@ def _compute_stats(
             # In the distributed case we stack multiple arrays to speed comms.
             if len(xs) > 1:
                 reduced_mus = jax.lax.pmean(
-                    jnp.stack(mus, axis=0), axis_name,
+                    jnp.stack(mus, axis=0),
+                    axis_name,
                     axis_index_groups=axis_index_groups,
                 )
                 return tuple(reduced_mus[i] for i in range(len(xs)))
             else:
-                return jax.lax.pmean(mus[0], axis_name, axis_index_groups=axis_index_groups)
+                return jax.lax.pmean(
+                    mus[0],
+                    axis_name,
+                    axis_index_groups=axis_index_groups
+                )
 
     if use_mean:
         if use_fast_variance:
@@ -176,7 +193,7 @@ def _normalize(
     x: ArrayLike,
     mean: Optional[ArrayLike],
     var: Optional[ArrayLike],
-    weights: Optional[ParamState],
+    weights: Optional[NormalizationParamState],
     reduction_axes: Axes,
     feature_axes: Axes,
     dtype: DTypeLike,
@@ -212,10 +229,9 @@ def _normalize(
         y = x - mean
         mul = jax.lax.rsqrt(var + epsilon)
         y = y * mul
-        args = []
         if weights is not None:
-            y, args = _scale_operation(y, weights.value)
-        dtype = canonicalize_dtype(x, *args, dtype=dtype)
+            y = weights.execute(y)
+        dtype = canonicalize_dtype(x, *jax.tree.leaves(weights.value), dtype=dtype)
     else:
         assert var is None, 'mean and val must be both None or not None.'
         assert weights is None, 'scale and bias are not supported without mean and val'
@@ -223,17 +239,6 @@ def _normalize(
     return jnp.asarray(y, dtype)
 
 
-def _scale_operation(x: jax.Array, param: Dict):
-    args = []
-    if 'scale' in param:
-        x = x * param['scale']
-        args.append(param['scale'])
-    if 'bias' in param:
-        x = x + param['bias']
-        args.append(param['bias'])
-    return x, args
-
-
 class _BatchNorm(Module):
     __module__ = 'brainstate.nn'
     num_spatial_dims: int
@@ -254,6 +259,8 @@ class _BatchNorm(Module):
         use_fast_variance: bool = True,
         name: Optional[str] = None,
         dtype: Any = None,
+        param_type: type = NormalizationParamState,
+        mean_type: type = BatchState,
     ):
         super().__init__(name=name)
 
@@ -279,8 +286,8 @@ class _BatchNorm(Module):
         feature_shape = tuple([(ax if i in self.feature_axes else 1)
                                for i, ax in enumerate(self.in_size)])
         if self.track_running_stats:
-            self.running_mean = LongTermState(jnp.zeros(feature_shape, dtype=self.dtype))
-            self.running_var = LongTermState(jnp.ones(feature_shape, dtype=self.dtype))
+            self.running_mean = mean_type(jnp.zeros(feature_shape, dtype=self.dtype))
+            self.running_var = mean_type(jnp.ones(feature_shape, dtype=self.dtype))
         else:
             self.running_mean = None
             self.running_var = None
@@ -290,7 +297,7 @@ class _BatchNorm(Module):
             assert track_running_stats, "Affine parameters are not needed when track_running_stats is False."
             bias = init.param(self.bias_initializer, feature_shape)
             scale = init.param(self.scale_initializer, feature_shape)
-            self.weight = ParamState(dict(bias=bias, scale=scale))
+            self.weight = param_type(dict(bias=bias, scale=scale))
         else:
             self.weight = None
 
@@ -531,6 +538,7 @@ class LayerNorm(Module):
         axis_index_groups: Any = None,
         use_fast_variance: bool = True,
         dtype: Optional[jax.typing.DTypeLike] = None,
+        param_type: type = NormalizationParamState,
     ):
         super().__init__()
 
@@ -554,7 +562,7 @@ class LayerNorm(Module):
         if use_bias:
             weights['bias'] = init.param(bias_init, feature_shape)
         if len(weights):
-            self.weight = ParamState(weights)
+            self.weight = param_type(weights)
         else:
             self.weight = None
 
@@ -654,6 +662,7 @@ class RMSNorm(Module):
         axis_name: Optional[str] = None,
         axis_index_groups: Any = None,
         use_fast_variance: bool = True,
+        param_type: type = NormalizationParamState,
     ):
         super().__init__()
 
@@ -663,7 +672,7 @@ class RMSNorm(Module):
         # parameters about axis
         feature_axes = (feature_axes,) if isinstance(feature_axes, int) else feature_axes
         self.feature_axes = _canonicalize_axes(len(self.in_size), feature_axes)
-        self.reduction_axes = (reduction_axes, ) if isinstance(reduction_axes, int) else reduction_axes
+        self.reduction_axes = (reduction_axes,) if isinstance(reduction_axes, int) else reduction_axes
         self.axis_name = axis_name
         self.axis_index_groups = axis_index_groups
 
@@ -671,7 +680,7 @@ class RMSNorm(Module):
         feature_shape = tuple([(ax if i in self.feature_axes else 1)
                                for i, ax in enumerate(self.in_size)])
         if use_scale:
-            self.scale = ParamState({'scale': init.param(scale_init, feature_shape)})
+            self.scale = param_type({'scale': init.param(scale_init, feature_shape)})
         else:
             self.scale = None
 
@@ -795,6 +804,7 @@ class GroupNorm(Module):
         axis_name: Optional[str] = None,
         axis_index_groups: Any = None,
         use_fast_variance: bool = True,
+        param_type: type = NormalizationParamState,
     ):
         super().__init__()
 
@@ -848,7 +858,7 @@ class GroupNorm(Module):
         if use_bias:
             weights['bias'] = init.param(bias_init, feature_shape)
         if len(weights):
-            self.weight = ParamState(weights)
+            self.weight = param_type(weights)
         else:
             self.weight = None
 

brainstate/util/_tracers.py

@@ -16,7 +16,6 @@ from __future__ import annotations
 
 import jax
 import jax.core
-from jax.interpreters import partial_eval as pe
 
 from ._pretty_repr import PrettyRepr, PrettyType, PrettyAttr
 
@@ -25,12 +24,6 @@ __all__ = [
 ]
 
 
-def new_jax_trace():
-    main = jax.core.thread_local_state.trace_state.trace_stack.stack[-1]
-    frame = main.jaxpr_stack[-1]
-    trace = pe.DynamicJaxprTrace(main, jax.core.cur_sublevel())
-    return frame, trace
-
 
 def current_jax_trace():
     """Returns the Jax tracing state."""

{brainstate-0.1.0.post20241129.dist-info → brainstate-0.1.0.post20241210.dist-info}/METADATA

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

{brainstate-0.1.0.post20241129.dist-info → brainstate-0.1.0.post20241210.dist-info}/RECORD

@@ -1,7 +1,5 @@
-benchmark/COBA_2005.py,sha256=Q8PsZ0lxu14jsF3bCtlZW35iQB8S2_oFEUYQzK2hPiA,5561
-benchmark/CUBA_2005.py,sha256=_W94yOMh2ueqblU4ItEPeTLwHF0_lbEWlVNEBy0Tix0,6222
 brainstate/__init__.py,sha256=r7C3eLTg8LEusoH6PGgBFFt4ZgbketYLoLA0lQhUCsE,2098
-brainstate/_state.py,sha256=0h3B32130Tvv41eeolYTbEGT9FZ9WUZ9yYoGyVnte_c,28808
+brainstate/_state.py,sha256=4aDpLyHGr1VlPXeLSfM3USQG5K4o7orF7IlaBdYrtfE,29098
 brainstate/_state_test.py,sha256=1boTp1w8DiCFLsPwNtlLrlIqGRpkasAmLid5bv2fgP4,2223
 brainstate/_utils.py,sha256=uJ6WWKq3yb05ZdktCQGLWOXsOJveL1H9pR7eev70Jes,1693
 brainstate/environ.py,sha256=G6r_rqfbofRbjFFalRu_DHaL7ruFTeLRXBQDXM6P-tQ,17477
@@ -17,24 +15,24 @@ brainstate/augment/_autograd_test.py,sha256=S2eEgrwTzdSi3u2nKE3u37WSThosLwx1WCP9
 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=bwoRm0KVK53gpB_b2dWwJl4BQplH74RCEU1VyE3DN4M,6762
+brainstate/augment/_mapping_test.py,sha256=TEAecjZmTSDCfxARgrzcDJ2dW1Yz_sCITmFiA9FGrhk,9455
 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_test.py,sha256=R1I76nG4zIqb6g3M_VxWts7rUC1OHJCjtQhPkcbXodk,1746
-brainstate/compile/_conditions.py,sha256=ocz6sDc7Xzabz2GnRsQmS6GDps-WP-OXUd0EZTTlG0k,10217
+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=SJmAfosVoGd4vhfFtb1IvjeFVW914bfTccCg6DoLWYk,1992
-brainstate/compile/_jit.py,sha256=bfEszNttEtE6npqHBam1_DBlRa39fE6qP6lGaWw2amA,13750
+brainstate/compile/_jit.py,sha256=3mQ-RUFz35wceZyKE_MoR58OBL0RK_i6sHm4rWYzMLs,13698
 brainstate/compile/_jit_test.py,sha256=zD7kck9SQJGmUDolh9P4luKwQ21fBGje1Z4STTEXIuA,4135
-brainstate/compile/_loop_collect_return.py,sha256=1IRhsWFad5khM1DIBLOI-kdC4pJYsT3TGqG0KO6ftjQ,22782
+brainstate/compile/_loop_collect_return.py,sha256=DybSBixeuxleKJV6n9FgVBDsUTmexzS0IdgWYRqp5cU,22940
 brainstate/compile/_loop_collect_return_test.py,sha256=bA-_11E8A_0jR5umEO3e409y7bb5QYDTgSL-SBaX7kQ,1802
-brainstate/compile/_loop_no_collection.py,sha256=2rSK20enkBMXPAbsCyb7PCICPNrgaSpl5jfumgWpxA0,7401
+brainstate/compile/_loop_no_collection.py,sha256=0i31gdQ7sI-d6pvnh08ttUUwdAtpx4uoYhGuf_CyL9s,7343
 brainstate/compile/_loop_no_collection_test.py,sha256=oStB1CSG_iLp9sHdXd1hJNFvlxbzjck9Iy4sABoJDj4,1419
-brainstate/compile/_make_jaxpr.py,sha256=9ZWYgW77M0YcrqmViyXSovuvQDoIa8hmSnhE7pBqDfU,32828
+brainstate/compile/_make_jaxpr.py,sha256=S5O9KUB3bsxoKcfptlV0MRfKA__Ija37WxkakIRL3z0,33010
 brainstate/compile/_make_jaxpr_test.py,sha256=qJUtkyj50JQ6f4UJbOLhvRdkbNn3NSKibFL9jESdQkA,4279
-brainstate/compile/_progress_bar.py,sha256=LML4DjrLSIeGYJWLjqy6BnHSz03fu1gnjf-7kljP384,3824
+brainstate/compile/_progress_bar.py,sha256=H544Oh10SiF5ccrKHM9ay7ZHigYIhNhSQGEKbDxRJgg,4485
 brainstate/compile/_unvmap.py,sha256=ewbLLNXiI_dBsEBaVzSS0BEXNol22sd9gMzk606lSkM,4139
 brainstate/compile/_util.py,sha256=aCvkTV--g4NsqcodTdBAISt4EwgezCbKzNUV58n-Q_Y,6304
 brainstate/event/__init__.py,sha256=wOBkq7kDg90M8Y9FuoXRlSEuu1ZzbIhCJ1dHeLqN6_Q,1194
@@ -52,7 +50,7 @@ brainstate/event/_misc.py,sha256=8IpPooXjF2m0-tuo3pGHqThq2yLSNmYziy_zdurZ3NI,104
 brainstate/event/_xla_custom_op.py,sha256=QB4jz_fUEPF-efJCVKAxwx8U79AqdcKoEg2QrGwot8I,10864
 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=tsCGoV_35IyZ0mM6_4fV9v0-Vj3V9Qm55U8wSpZ_E4o,22180
+brainstate/functional/_activations.py,sha256=S0Ok7sq5FTbmJWSejpOCHo1jpKX0gYOLy_TO2IUXM8s,21726
 brainstate/functional/_activations_test.py,sha256=T___RlSrIfXwlkw8dg5A9EZMTZGDzv3a2evUwq_nYFg,13034
 brainstate/functional/_normalization.py,sha256=i2EV7hSsqcNdcYRX2wAxjq8doHwyN9eNJTGTaPt03xE,2605
 brainstate/functional/_others.py,sha256=_u_Ys-LiLzDAP4zJggVwaVvirgoS3jvhXMREoS6JOkM,1737
@@ -63,7 +61,7 @@ brainstate/graph/_graph_context_test.py,sha256=IYpjqbXwSFF65XL0ZbdPeC1jYyEHLpQVr
 brainstate/graph/_graph_convert.py,sha256=llSREtGQrIggkD0wmxUbYKuSveLW4ihDZME6Ab-mRTQ,9147
 brainstate/graph/_graph_node.py,sha256=BTuVlGgA2b82zNudjsN88QXuxfDcMvU2-kB64AkdQnY,8993
 brainstate/graph/_graph_node_test.py,sha256=BFGfdzZFDHI0XK7hHotSVWKt3em1taGvn8FHF9NCXx8,2702
-brainstate/graph/_graph_operation.py,sha256=C_WDUwg_cY7mduW8lb8yRG3trJy9KCaRmMluXf003Ss,64006
+brainstate/graph/_graph_operation.py,sha256=PupZeFWBR-OHbhdJcoqlvy2YqoIS9Ze4q0tz8HRy4f4,64166
 brainstate/graph/_graph_operation_test.py,sha256=ADyyuMk2xidEkkFNpGvUbvEtRmUj-tqOI4cF3eRuakM,24678
 brainstate/init/__init__.py,sha256=R1dHgub47o-WJM9QkFLc7x_Q7GsyaKKDtrRHTFPpC5g,1097
 brainstate/init/_base.py,sha256=B_NLS9aKNrvuj5NAlSgBbQTVev7IRvzcx8vH0J-Gq2w,1671
@@ -73,7 +71,7 @@ brainstate/init/_random_inits_test.py,sha256=lBL2RQdBSZ88Zqz4IMdbHJMvDi7ooZq6caC
 brainstate/init/_regular_inits.py,sha256=DmVMajugfyYFNUMzgFdDKMvbBu9hMWxkfDd-50uhoLg,3187
 brainstate/init/_regular_inits_test.py,sha256=tJl4aOkclllJIfKzJTbc0cfYCw2SoBsx8_G123RnqbU,1842
 brainstate/nn/__init__.py,sha256=rxURT8J1XfBn3Vh3Dx_WzVADWn9zVriIty5KZEG-x6o,1622
-brainstate/nn/_collective_ops.py,sha256=BzKKUxeBPtZojHqvOzgPkz9EFaVJUwQWmqR18mF2l38,6233
+brainstate/nn/_collective_ops.py,sha256=sSjIIs1MvZA30XFFmK7iL1D_sCeh7hFd3PanCH6kgZo,6779
 brainstate/nn/_exp_euler.py,sha256=yjkfSllFxGWKEAlHo5AzBizzkFj6FEVDKmFV6E2g214,3521
 brainstate/nn/_exp_euler_test.py,sha256=clwRD8QR71k1jn6NrACMDEUcFMh0J9RTosoPnlYWUkw,1242
 brainstate/nn/_module.py,sha256=HDLPvLfB7jat2VT3gBu0MxA7vfzK7xgowemitHX8Cgo,10835
@@ -82,7 +80,7 @@ brainstate/nn/metrics.py,sha256=iupHjSRTHYY-HmEPBC4tXWrZfF4zh1ek2NwSAA0gnwE,1473
 brainstate/nn/_dyn_impl/__init__.py,sha256=Oazar7h89dp1WA2Vx4Tj7gCBhxJKH4LAUEABkBEG7vU,1462
 brainstate/nn/_dyn_impl/_dynamics_neuron.py,sha256=cTbIn41EPYG0h3ICzKBXxpgB6wwA2K8k5FAcf3Pa5N8,10927
 brainstate/nn/_dyn_impl/_dynamics_neuron_test.py,sha256=Tfzrzu7udGrLJGnqItiLWe5WT0dgduvYOgzGCnaPJQg,6317
-brainstate/nn/_dyn_impl/_dynamics_synapse.py,sha256=BP-ko0FyjWZopuUhAy3Ot3wWRQlGcpumWJpKrQakqok,11869
+brainstate/nn/_dyn_impl/_dynamics_synapse.py,sha256=IHy6IsGjWpKZ8NLq4X7PaRwx3tpO2HRZNppCWM2fe4I,11862
 brainstate/nn/_dyn_impl/_dynamics_synapse_test.py,sha256=t5i-HV0ii9sUNzWTEv04o26QVtQ-mCdMJcFq2MD755A,4981
 brainstate/nn/_dyn_impl/_inputs.py,sha256=6eZKnkmrM0Gog2fpSKjSnwnQvhbFYhG4q9Vuo-GH2LI,5050
 brainstate/nn/_dyn_impl/_projection_alignpost.py,sha256=PNC1Tzx_SF2DHAHeJCufXzO_Q4qLoBpWABI45B3GRuc,876
@@ -98,7 +96,7 @@ brainstate/nn/_dynamics/_state_delay.py,sha256=nZYGmVKmQvAQu-W4YOUFH1gnr-ZS3rg_G
 brainstate/nn/_dynamics/_synouts.py,sha256=9TGAc-nVa50th7KKn4oKLbro-4W4rwxYvp-eu7ksAIE,4491
 brainstate/nn/_dynamics/_synouts_test.py,sha256=V_jDswRN4VvEXD-2yJO3VA1TALgX0HK6oPBQiUntOWc,2266
 brainstate/nn/_elementwise/__init__.py,sha256=PK8oq1K_EG2941AiUyLxCWoRdWvMO3yt8ZJbw3Lkhu8,935
-brainstate/nn/_elementwise/_dropout.py,sha256=lXy8ki5mnmhU9lacJGesy7mums0qH1pplXD6BGGdw-I,17338
+brainstate/nn/_elementwise/_dropout.py,sha256=0Ebo-2y1VswvBqZ7sCA0SEUm37y49EUsef8oiSFpYGk,17759
 brainstate/nn/_elementwise/_dropout_test.py,sha256=Qn7xqZOyZMPCGF6tFjTiPId0yELOXjSsW5-hgihP3fE,4383
 brainstate/nn/_elementwise/_elementwise.py,sha256=om-KpwDTk5yFG5KBYXXHquRLV7s28_FJjk-omvyMyvQ,33342
 brainstate/nn/_elementwise/_elementwise_test.py,sha256=SZI9jB39sZ5SO1dpWGW-PhodthwN0GU9FY1nqf2fWcs,5341
@@ -108,7 +106,7 @@ brainstate/nn/_interaction/_conv_test.py,sha256=fHXRFYnDghFiKre63RqMwIE_gbPKdK34
 brainstate/nn/_interaction/_embedding.py,sha256=iK0I1ExKWFa_QzV9UDGj32Ljsmdr1g_LlAtMcusebxU,2187
 brainstate/nn/_interaction/_linear.py,sha256=bjiWGJCe81ugQQOykwjWlLW5uhe0CHWwkPA20a4n5YQ,21340
 brainstate/nn/_interaction/_linear_test.py,sha256=KlvFZA0rpyaspf4LT4K7u-RR5jCEB_q1WReqAw9sFcU,1274
-brainstate/nn/_interaction/_normalizations.py,sha256=pdVzuMMmKP49pFbLxZJ6A9m0_nBEMEvOUN4Y1lWnhf0,37004
+brainstate/nn/_interaction/_normalizations.py,sha256=7YDzkmO_iqd70fH_wawb60Bu8eGOdvZq23emP-b68Hc,37440
 brainstate/nn/_interaction/_normalizations_test.py,sha256=2p1Jf8nA999VYGWbvOZfKYlKk6UmL0vaEB76xkXxkXw,2438
 brainstate/nn/_interaction/_poolings.py,sha256=LpwuyeNBVCaVFW7zWc7E-vvlYqx54h46Br5XT6zd_94,47020
 brainstate/nn/_interaction/_poolings_test.py,sha256=wmd5PngZ3E9tNyF3s0xk-DoDR5yFqpTi9A6nbNoIqn4,7429
@@ -136,10 +134,10 @@ brainstate/util/_others.py,sha256=jsPZwP-v_5HRV-LB5F0NUsiqr04y8bmGIsu_JMyVcbQ,14
 brainstate/util/_pretty_repr.py,sha256=NYEBCo2iz9Potx-IR7uZZzt2aLQW_94vH79fGusiC2A,5737
 brainstate/util/_scaling.py,sha256=pc_eM_SZVwkY65I4tJh1ODiHNCoEhsfFXl2zBK0PLAg,7562
 brainstate/util/_struct.py,sha256=0exv0oOiSt1hmx20Y4J2-pCGtCTx13WcAlEYSBkyung,17640
-brainstate/util/_tracers.py,sha256=-sX76GJRThdSpDJBejAIzDdBbVhmH6kb-1WoDJVI7V0,2556
+brainstate/util/_tracers.py,sha256=0r5T4nhxMzI79NtqroqitsdMT4YfpgV5RdYJLS5uJ0w,2285
 brainstate/util/_visualization.py,sha256=n4ZVz10z7VBqA0cKO6vyHwEMprWJgPeEqtITzDMai2Y,1519
-brainstate-0.1.0.post20241129.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.1.0.post20241129.dist-info/METADATA,sha256=s2ks57prKGe9kPi_goPGwYVGr40TrgrEFjtC0W8ff6Q,3401
-brainstate-0.1.0.post20241129.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.1.0.post20241129.dist-info/top_level.txt,sha256=MVkn5SZk0qis32_AGxU2Tg1xADfj2IgCNS25CQD7_ng,21
-brainstate-0.1.0.post20241129.dist-info/RECORD,,
+brainstate-0.1.0.post20241210.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.1.0.post20241210.dist-info/METADATA,sha256=E6AATarjpwssXflLfA-OCkxFxqZxqJNxHZteO6UWMhw,3401
+brainstate-0.1.0.post20241210.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+brainstate-0.1.0.post20241210.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.1.0.post20241210.dist-info/RECORD,,

{brainstate-0.1.0.post20241129.dist-info → brainstate-0.1.0.post20241210.dist-info}/top_level.txt

@@ -1,2 +1 @@
-benchmark
 brainstate

benchmark/COBA_2005.py

@@ -1,125 +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.
-# ==============================================================================
-
-#
-# Implementation of the paper:
-#
-# - Brette, R., Rudolph, M., Carnevale, T., Hines, M., Beeman, D., Bower, J. M., et al. (2007),
-#   Simulation of networks of spiking neurons: a review of tools and strategies., J. Comput. Neurosci., 23, 3, 349–98
-#
-# which is based on the balanced network proposed by:
-#
-# - Vogels, T. P. and Abbott, L. F. (2005), Signal propagation and logic gating in networks of integrate-and-fire neurons., J. Neurosci., 25, 46, 10786–95
-#
-import os
-import sys
-
-sys.path.append('../')
-os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '0.99'
-os.environ['JAX_TRACEBACK_FILTERING'] = 'off'
-
-
-import jax
-import brainunit as u
-import time
-import brainstate as bst
-
-
-class EINet(bst.nn.DynamicsGroup):
-    def __init__(self, scale):
-        super().__init__()
-        self.n_exc = int(3200 * scale)
-        self.n_inh = int(800 * scale)
-        self.num = self.n_exc + self.n_inh
-        self.N = bst.nn.LIFRef(self.num, V_rest=-60. * u.mV, V_th=-50. * u.mV, V_reset=-60. * u.mV,
-                               tau=20. * u.ms, tau_ref=5. * u.ms,
-                               V_initializer=bst.init.Normal(-55., 2., unit=u.mV))
-        self.E = bst.nn.AlignPostProj(
-            comm=bst.event.FixedProb(self.n_exc, self.num, prob=80 / self.num, weight=0.6 * u.mS),
-            syn=bst.nn.Expon.desc(self.num, tau=5. * u.ms),
-            out=bst.nn.COBA.desc(E=0. * u.mV),
-            post=self.N
-        )
-        self.I = bst.nn.AlignPostProj(
-            comm=bst.event.FixedProb(self.n_inh, self.num, prob=80 / self.num, weight=6.7 * u.mS),
-            syn=bst.nn.Expon.desc(self.num, tau=10. * u.ms),
-            out=bst.nn.COBA.desc(E=-80. * u.mV),
-            post=self.N
-        )
-
-    def init_state(self, *args, **kwargs):
-        self.rate = bst.ShortTermState(u.math.zeros(self.num))
-
-    def update(self, t, inp):
-        with bst.environ.context(t=t):
-            spk = self.N.get_spike() != 0.
-            self.E(spk[:self.n_exc])
-            self.I(spk[self.n_exc:])
-            self.N(inp)
-            self.rate.value += self.N.get_spike()
-
-
-@bst.compile.jit(static_argnums=0)
-def run(scale: float):
-    # network
-    net = EINet(scale)
-    bst.nn.init_all_states(net)
-
-    duration = 1e4 * u.ms
-    # simulation
-    with bst.environ.context(dt=0.1 * u.ms):
-        times = u.math.arange(0. * u.ms, duration, bst.environ.get_dt())
-        bst.compile.for_loop(lambda t: net.update(t, 20. * u.mA), times)
-
-    return net.num, net.rate.value.sum() / net.num / duration.to_decimal(u.second)
-
-
-for s in [1, 2, 4, 6, 8, 10, 20, 40, 60, 80, 100]:
-    jax.block_until_ready(run(s))
-
-    t0 = time.time()
-    n, rate = jax.block_until_ready(run(s))
-    t1 = time.time()
-    print(f'scale={s}, size={n}, time = {t1 - t0} s, firing rate = {rate} Hz')
-
-
-# A6000 NVIDIA GPU
-
-# scale=1, size=4000, time = 2.659956455230713 s, firing rate = 50.62445068359375 Hz
-# scale=2, size=8000, time = 2.7318649291992188 s, firing rate = 50.613040924072266 Hz
-# scale=4, size=16000, time = 2.807222604751587 s, firing rate = 50.60573959350586 Hz
-# scale=6, size=24000, time = 3.026782512664795 s, firing rate = 50.60918045043945 Hz
-# scale=8, size=32000, time = 3.1258811950683594 s, firing rate = 50.607574462890625 Hz
-# scale=10, size=40000, time = 3.172346353530884 s, firing rate = 50.60942840576172 Hz
-# scale=20, size=80000, time = 3.751189947128296 s, firing rate = 50.612369537353516 Hz
-# scale=40, size=160000, time = 5.0217814445495605 s, firing rate = 50.617958068847656 Hz
-# scale=60, size=240000, time = 7.002646207809448 s, firing rate = 50.61948776245117 Hz
-# scale=80, size=320000, time = 9.384576320648193 s, firing rate = 50.618499755859375 Hz
-# scale=100, size=400000, time = 11.69654369354248 s, firing rate = 50.61605453491211 Hz
-
-
-# AMD Ryzen 7 7840HS
-
-# scale=1, size=4000, time = 4.436027526855469 s, firing rate = 50.6119270324707 Hz
-# scale=2, size=8000, time = 8.349745273590088 s, firing rate = 50.612266540527344 Hz
-# scale=4, size=16000, time = 16.39163303375244 s, firing rate = 50.61349105834961 Hz
-# scale=6, size=24000, time = 15.725558042526245 s, firing rate = 50.6125602722168 Hz
-# scale=8, size=32000, time = 21.31995177268982 s, firing rate = 50.61244583129883 Hz
-# scale=10, size=40000, time = 27.811061143875122 s, firing rate = 50.61423873901367 Hz
-# scale=20, size=80000, time = 45.54235219955444 s, firing rate = 50.61320877075195 Hz
-# scale=40, size=160000, time = 82.22228026390076 s, firing rate = 50.61309814453125 Hz
-# scale=60, size=240000, time = 125.44037556648254 s, firing rate = 50.613094329833984 Hz
-# scale=80, size=320000, time = 171.20458459854126 s, firing rate = 50.613365173339844 Hz
-# scale=100, size=400000, time = 215.4547393321991 s, firing rate = 50.6129150390625 Hz

benchmark/CUBA_2005.py

@@ -1,149 +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.
-# ==============================================================================
-
-#
-# Implementation of the paper:
-#
-# - Brette, R., Rudolph, M., Carnevale, T., Hines, M., Beeman, D., Bower, J. M., et al. (2007),
-#   Simulation of networks of spiking neurons: a review of tools and strategies., J. Comput. Neurosci., 23, 3, 349–98
-#
-# which is based on the balanced network proposed by:
-#
-# - Vogels, T. P. and Abbott, L. F. (2005), Signal propagation and logic gating in networks of integrate-and-fire neurons., J. Neurosci., 25, 46, 10786–95
-#
-
-import os
-import sys
-
-sys.path.append('../')
-os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '0.99'
-os.environ['JAX_TRACEBACK_FILTERING'] = 'off'
-
-
-import jax
-import time
-
-import brainunit as u
-
-import brainstate as bst
-
-
-
-class FixedProb(bst.nn.Module):
-    def __init__(self, n_pre, n_post, prob, weight):
-        super().__init__()
-        self.prob = prob
-        self.weight = weight
-        self.n_pre = n_pre
-        self.n_post = n_post
-
-        self.mask = bst.random.rand(n_pre, n_post) < prob
-
-    def update(self, x):
-        return (x @ self.mask) * self.weight
-
-
-class EINet(bst.nn.DynamicsGroup):
-    def __init__(self, scale=1.0):
-        super().__init__()
-        self.n_exc = int(3200 * scale)
-        self.n_inh = int(800 * scale)
-        self.num = self.n_exc + self.n_inh
-        self.N = bst.nn.LIFRef(
-            self.num, V_rest=-49. * u.mV, V_th=-50. * u.mV, V_reset=-60. * u.mV,
-            tau=20. * u.ms, tau_ref=5. * u.ms,
-            V_initializer=bst.init.Normal(-55., 2., unit=u.mV)
-        )
-        self.E = bst.nn.AlignPostProj(
-            comm=bst.event.FixedProb(self.n_exc, self.num, prob=80 / self.num, weight=1.62 * u.mS),
-            # comm=FixedProb(self.n_exc, self.num, prob=80 / self.num, weight=1.62 * u.mS),
-            syn=bst.nn.Expon.desc(self.num, tau=5. * u.ms),
-            out=bst.nn.CUBA.desc(scale=u.volt),
-            post=self.N
-        )
-        self.I = bst.nn.AlignPostProj(
-            comm=bst.event.FixedProb(self.n_inh, self.num, prob=80 / self.num, weight=-9.0 * u.mS),
-            # comm=FixedProb(self.n_inh, self.num, prob=80 / self.num, weight=-9.0 * u.mS),
-            syn=bst.nn.Expon.desc(self.num, tau=10. * u.ms),
-            out=bst.nn.CUBA.desc(scale=u.volt),
-            post=self.N
-        )
-
-    def init_state(self, *args, **kwargs):
-        self.rate = bst.ShortTermState(u.math.zeros(self.num))
-
-    def update(self, t, inp):
-        with bst.environ.context(t=t):
-            spk = self.N.get_spike()
-            self.E(spk[:self.n_exc])
-            self.I(spk[self.n_exc:])
-            self.N(inp)
-            self.rate.value += self.N.get_spike()
-
-
-@bst.compile.jit(static_argnums=0)
-def run(scale: float):
-    # network
-    net = EINet(scale)
-    bst.nn.init_all_states(net)
-
-    duration = 1e4 * u.ms
-    # simulation
-    with bst.environ.context(dt=0.1 * u.ms):
-        times = u.math.arange(0. * u.ms, duration, bst.environ.get_dt())
-        bst.compile.for_loop(lambda t: net.update(t, 20. * u.mA), times,
-                             # pbar=bst.compile.ProgressBar(100)
-                             )
-
-    return net.num, net.rate.value.sum() / net.num / duration.to_decimal(u.second)
-
-
-for s in [1, 2, 4, 6, 8, 10, 20, 40, 60, 80, 100]:
-    jax.block_until_ready(run(s))
-
-    t0 = time.time()
-    n, rate = jax.block_until_ready(run(s))
-    t1 = time.time()
-    print(f'scale={s}, size={n}, time = {t1 - t0} s, firing rate = {rate} Hz')
-
-
-# A6000 NVIDIA GPU
-
-# scale=1, size=4000, time = 2.6354849338531494 s, firing rate = 24.982027053833008 Hz
-# scale=2, size=8000, time = 2.6781561374664307 s, firing rate = 23.719463348388672 Hz
-# scale=4, size=16000, time = 2.7448785305023193 s, firing rate = 24.592931747436523 Hz
-# scale=6, size=24000, time = 2.8237478733062744 s, firing rate = 24.159996032714844 Hz
-# scale=8, size=32000, time = 2.9344418048858643 s, firing rate = 24.956790924072266 Hz
-# scale=10, size=40000, time = 3.042517900466919 s, firing rate = 23.644424438476562 Hz
-# scale=20, size=80000, time = 3.6727631092071533 s, firing rate = 24.226743698120117 Hz
-# scale=40, size=160000, time = 4.857396602630615 s, firing rate = 24.329742431640625 Hz
-# scale=60, size=240000, time = 6.812030792236328 s, firing rate = 24.370006561279297 Hz
-# scale=80, size=320000, time = 9.227966547012329 s, firing rate = 24.41067886352539 Hz
-# scale=100, size=400000, time = 11.405697584152222 s, firing rate = 24.32524871826172 Hz
-
-
-# AMD Ryzen 7 7840HS
-
-# scale=1, size=4000, time = 1.1661601066589355 s, firing rate = 22.438201904296875 Hz
-# scale=2, size=8000, time = 3.3255884647369385 s, firing rate = 23.868364334106445 Hz
-# scale=4, size=16000, time = 6.950139999389648 s, firing rate = 24.21693229675293 Hz
-# scale=6, size=24000, time = 10.011993169784546 s, firing rate = 24.240270614624023 Hz
-# scale=8, size=32000, time = 13.027734518051147 s, firing rate = 24.753198623657227 Hz
-# scale=10, size=40000, time = 16.449942350387573 s, firing rate = 24.7176570892334 Hz
-# scale=20, size=80000, time = 30.754598140716553 s, firing rate = 24.119956970214844 Hz
-# scale=40, size=160000, time = 63.6387836933136 s, firing rate = 24.72784996032715 Hz
-# scale=60, size=240000, time = 78.58532166481018 s, firing rate = 24.402742385864258 Hz
-# scale=80, size=320000, time = 102.4250214099884 s, firing rate = 24.59092140197754 Hz
-# scale=100, size=400000, time = 145.35173273086548 s, firing rate = 24.33751106262207 Hz