brainstate 0.1.0.post20241125__py2.py3-none-any.whl → 0.1.0.post20241209__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/_loop_collect_return.py

@@ -211,7 +211,11 @@ def scan(
 
     # 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/_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,19 +93,37 @@ 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.lax.cond(
+        jax.debug.callback(
+            self._tqdm,
             iter_num == 0,
-            lambda: jax.debug.callback(self._define_tqdm),
-            lambda: None,
-        )
-        _ = jax.lax.cond(
             (iter_num + 1) % self.print_freq == 0,
-            lambda: jax.debug.callback(self._update_tqdm),
-            lambda: None,
-        )
-        _ = jax.lax.cond(
-            iter_num == self.n - 1,
-            lambda: jax.debug.callback(self._close_tqdm),
-            lambda: None,
+            iter_num == self.n - 1
         )
+
+        # _ = jax.lax.cond(
+        #     iter_num == 0,
+        #     lambda: jax.debug.callback(self._define_tqdm, ordered=True),
+        #     lambda: None,
+        # )
+        # _ = jax.lax.cond(
+        #     (iter_num + 1) % self.print_freq == 0,
+        #     lambda: jax.debug.callback(self._update_tqdm, ordered=True),
+        #     lambda: None,
+        # )
+        # _ = jax.lax.cond(
+        #     iter_num == self.n - 1,
+        #     lambda: jax.debug.callback(self._close_tqdm, ordered=True),
+        #     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
 
@@ -88,12 +95,11 @@ class _DropoutNd(ElementWiseBlock):
         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}."
+        assert 0. <= prob <= 1., f"Dropout probability must be in the range [0, 1]. But got {prob}."
         self.prob = prob
         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):
@@ -112,12 +118,15 @@ class _DropoutNd(ElementWiseBlock):
         fit_phase = environ.get('fit', desc='Whether this is a fitting process. Bool.')
 
         # generate mask
-        if fit_phase:
+        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):
@@ -396,7 +405,7 @@ class DropoutFixed(ElementWiseBlock):
         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}."
+        assert 0. <= prob <= 1., f"Dropout probability must be in the range [0, 1]. But got {prob}."
         self.prob = prob
         self.in_size = in_size
         self.out_size = in_size
@@ -407,7 +416,7 @@ class DropoutFixed(ElementWiseBlock):
     def update(self, x):
         dtype = u.math.get_dtype(x)
         fit_phase = environ.get('fit', desc='Whether this is a fitting process. Bool.')
-        if fit_phase:
+        if fit_phase and self.prob < 1.:
             if self.mask.value.shape != x.shape:
                 raise ValueError(f"Input shape {x.shape} does not match the mask shape {self.mask.value.shape}. "
                                  f"Please call `init_state()` method first.")