brainstate 0.1.3__py2.py3-none-any.whl → 0.1.4__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.3"
+__version__ = "0.1.4"
 
 from . import augment
 from . import compile

brainstate/_compatible_import.py

@@ -16,10 +16,9 @@
 # -*- coding: utf-8 -*-
 
 
-import importlib.util
 from contextlib import contextmanager
 from functools import partial
-from typing import Iterable, Hashable, TypeVar, Callable, TYPE_CHECKING
+from typing import Iterable, Hashable, TypeVar, Callable
 
 import jax
 
@@ -31,7 +30,6 @@ __all__ = [
     'get_aval',
     'Tracer',
     'to_concrete_aval',
-    'brainevent',
     'safe_map',
     'safe_zip',
     'unzip2',
@@ -47,8 +45,6 @@ T3 = TypeVar("T3")
 
 from saiunit._compatible_import import wrap_init
 
-brainevent_installed = importlib.util.find_spec('brainevent') is not None
-
 from jax.core import get_aval, Tracer
 
 if jax.__version_info__ < (0, 5, 0):
@@ -150,14 +146,3 @@ def to_concrete_aval(aval):
         return aval.to_concrete_value()
     return aval
 
-
-if not brainevent_installed:
-    if not TYPE_CHECKING:
-        class BrainEvent:
-            def __getattr__(self, item):
-                raise ImportError('brainevent is not installed, please install brainevent first.')
-
-        brainevent = BrainEvent()
-
-else:
-    import brainevent

brainstate/compile/_jit.py

@@ -51,6 +51,7 @@ def _get_jitted_fun(
     out_shardings,
     static_argnums,
     donate_argnums,
+    static_argnames,
     donate_argnames,
     keep_unused,
     device,
@@ -59,10 +60,12 @@ def _get_jitted_fun(
     abstracted_axes,
     **kwargs
 ) -> JittedFunction:
-    static_argnums = _ensure_index_tuple(tuple() if static_argnums is None else static_argnums)
+    static_argnums = tuple() if static_argnums is None else _ensure_index_tuple(static_argnums)
+    donate_argnums = tuple() if donate_argnums is None else _ensure_index_tuple(donate_argnums)
     fun = StatefulFunction(
         fun,
         static_argnums=static_argnums,
+        static_argnames=static_argnames,
         abstracted_axes=abstracted_axes,
         cache_type='jit',
         name='jit'
@@ -70,7 +73,8 @@ def _get_jitted_fun(
     jit_fun = jax.jit(
         fun.jaxpr_call,
         static_argnums=tuple(i + 1 for i in static_argnums),
-        donate_argnums=donate_argnums,
+        static_argnames=static_argnames,
+        donate_argnums=tuple(i + 1 for i in donate_argnums),
         donate_argnames=donate_argnames,
         keep_unused=keep_unused,
         device=device,
@@ -179,6 +183,7 @@ def jit(
     out_shardings=sharding_impls.UNSPECIFIED,
     static_argnums: int | Sequence[int] | None = None,
     donate_argnums: int | Sequence[int] | None = None,
+    static_argnames: str | Sequence[str] | None = None,
     donate_argnames: str | Iterable[str] | None = None,
     keep_unused: bool = False,
     device: Device | None = None,
@@ -190,9 +195,6 @@ def jit(
     """
     Sets up ``fun`` for just-in-time compilation with XLA.
 
-    Does not support setting ``static_argnames`` as in ``jax.jit()``.
-
-
     Args:
       fun: Function to be jitted.
       in_shardings: Pytree of structure matching that of arguments to ``fun``,
@@ -246,6 +248,11 @@ def jit(
         provided, ``inspect.signature`` is not used, and only actual
         parameters listed in either ``static_argnums`` or ``static_argnames`` will
         be treated as static.
+      static_argnames: An optional string or collection of strings specifying
+        which named arguments are treated as static (compile-time constant).
+        Operations that only depend on static arguments will be constant-folded in
+        Python (during tracing), and so the corresponding argument values can be
+        any Python object.
       donate_argnums: Specify which positional argument buffers are "donated" to
         the computation. It is safe to donate argument buffers if you no longer
         need them once the computation has finished. In some cases XLA can make
@@ -309,6 +316,7 @@ def jit(
                 out_shardings=out_shardings,
                 static_argnums=static_argnums,
                 donate_argnums=donate_argnums,
+                static_argnames=static_argnames,
                 donate_argnames=donate_argnames,
                 keep_unused=keep_unused,
                 device=device,
@@ -327,6 +335,7 @@ def jit(
             out_shardings,
             static_argnums,
             donate_argnums,
+            static_argnames,
             donate_argnames,
             keep_unused,
             device,

brainstate/compile/_make_jaxpr.py

@@ -88,6 +88,12 @@ __all__ = [
 ]
 
 
+def _ensure_str(x: str) -> str:
+    if not isinstance(x, str):
+        raise TypeError(f"argument is not a string: {x}")
+    return x
+
+
 def _ensure_index_tuple(x: Any) -> tuple[int, ...]:
     """Convert x to a tuple of indices."""
     x = jax.core.concrete_or_error(None, x, "expected a static index or sequence of indices.")
@@ -97,6 +103,14 @@ def _ensure_index_tuple(x: Any) -> tuple[int, ...]:
         return tuple(safe_map(operator.index, x))
 
 
+def _ensure_str_tuple(x: str | Iterable[str]) -> tuple[str, ...]:
+    """Convert x to a tuple of strings."""
+    if isinstance(x, str):
+        return (x,)
+    else:
+        return tuple(safe_map(_ensure_str, x))
+
+
 def _jax_v04_new_arg_fn(frame, trace, aval):
     """
     Transform a new argument to a tracer.
@@ -155,6 +169,9 @@ def _init_state_trace_stack(name) -> StateTraceStack:
     return state_trace
 
 
+default_cache_key = ((), ())
+
+
 class StatefulFunction(PrettyObject):
     """
     A wrapper class for a function that collects the states that are read and written by the function. The states are
@@ -170,6 +187,7 @@ class StatefulFunction(PrettyObject):
         arguments and return value should be arrays, scalars, or standard Python
         containers (tuple/list/dict) thereof.
       static_argnums: See the :py:func:`jax.jit` docstring.
+      static_argnames: See the :py:func:`jax.jit` docstring.
       axis_env: Optional, a sequence of pairs where the first element is an axis
           name and the second element is a positive integer representing the size of
           the mapped axis with that name. This parameter is useful when lowering
@@ -199,6 +217,7 @@ class StatefulFunction(PrettyObject):
         self,
         fun: Callable,
         static_argnums: Union[int, Iterable[int]] = (),
+        static_argnames: Union[str, Iterable[str]] = (),
         axis_env: Optional[Sequence[tuple[Hashable, int]]] = None,
         abstracted_axes: Optional[Any] = None,
         state_returns: Union[str, Tuple[str, ...]] = ('read', 'write'),
@@ -207,11 +226,12 @@ class StatefulFunction(PrettyObject):
     ):
         # explicit parameters
         self.fun = fun
-        self.static_argnums = _ensure_index_tuple(tuple() if static_argnums is None else static_argnums)
+        self.static_argnums = tuple() if static_argnums is None else _ensure_index_tuple(static_argnums)
+        self.static_argnames = tuple() if static_argnames is None else _ensure_str_tuple(static_argnames)
         self.axis_env = axis_env
         self.abstracted_axes = abstracted_axes
         self.state_returns = tuple(state_returns) if isinstance(state_returns, (tuple, list)) else (state_returns,)
-        assert cache_type in [None, 'jit']
+        assert cache_type in [None, 'jit'], f"Invalid cache type: {cache_type}"
         self.name = name
 
         # implicit parameters
@@ -226,7 +246,7 @@ class StatefulFunction(PrettyObject):
             return None
         return k, v
 
-    def get_jaxpr(self, cache_key: Hashable = ()) -> ClosedJaxpr:
+    def get_jaxpr(self, cache_key: Hashable = None) -> ClosedJaxpr:
         """
         Read the JAX Jaxpr representation of the function.
 
@@ -236,11 +256,13 @@ class StatefulFunction(PrettyObject):
         Returns:
           The JAX Jaxpr representation of the function.
         """
+        if cache_key is None:
+            cache_key = default_cache_key
         if cache_key not in self._cached_jaxpr:
             raise ValueError(f"the function is not called with the static arguments: {cache_key}")
         return self._cached_jaxpr[cache_key]
 
-    def get_out_shapes(self, cache_key: Hashable = ()) -> PyTree:
+    def get_out_shapes(self, cache_key: Hashable = None) -> PyTree:
         """
         Read the output shapes of the function.
 
@@ -250,11 +272,13 @@ class StatefulFunction(PrettyObject):
         Returns:
           The output shapes of the function.
         """
+        if cache_key is None:
+            cache_key = default_cache_key
         if cache_key not in self._cached_out_shapes:
             raise ValueError(f"the function is not called with the static arguments: {cache_key}")
         return self._cached_out_shapes[cache_key]
 
-    def get_out_treedef(self, cache_key: Hashable = ()) -> PyTree:
+    def get_out_treedef(self, cache_key: Hashable = None) -> PyTree:
         """
         Read the output tree of the function.
 
@@ -264,11 +288,13 @@ class StatefulFunction(PrettyObject):
         Returns:
           The output tree of the function.
         """
+        if cache_key is None:
+            cache_key = default_cache_key
         if cache_key not in self._cached_jaxpr_out_tree:
             raise ValueError(f"the function is not called with the static arguments: {cache_key}")
         return self._cached_jaxpr_out_tree[cache_key]
 
-    def get_state_trace(self, cache_key: Hashable = ()) -> StateTraceStack:
+    def get_state_trace(self, cache_key: Hashable = None) -> StateTraceStack:
         """
         Read the state trace of the function.
 
@@ -278,11 +304,13 @@ class StatefulFunction(PrettyObject):
         Returns:
           The state trace of the function.
         """
+        if cache_key is None:
+            cache_key = default_cache_key
         if cache_key not in self._cached_state_trace:
             raise ValueError(f"the function is not called with the static arguments: {cache_key}")
         return self._cached_state_trace[cache_key]
 
-    def get_states(self, cache_key: Hashable = ()) -> Tuple[State, ...]:
+    def get_states(self, cache_key: Hashable = None) -> Tuple[State, ...]:
         """
         Read the states that are read and written by the function.
 
@@ -292,9 +320,11 @@ class StatefulFunction(PrettyObject):
         Returns:
           The states that are read and written by the function.
         """
+        if cache_key is None:
+            cache_key = default_cache_key
         return tuple(self.get_state_trace(cache_key).states)
 
-    def get_read_states(self, cache_key: Hashable = ()) -> Tuple[State, ...]:
+    def get_read_states(self, cache_key: Hashable = None) -> Tuple[State, ...]:
         """
         Read the states that are read by the function.
 
@@ -304,9 +334,11 @@ class StatefulFunction(PrettyObject):
         Returns:
           The states that are read by the function.
         """
+        if cache_key is None:
+            cache_key = default_cache_key
         return self.get_state_trace(cache_key).get_read_states()
 
-    def get_write_states(self, cache_key: Hashable = ()) -> Tuple[State, ...]:
+    def get_write_states(self, cache_key: Hashable = None) -> Tuple[State, ...]:
         """
         Read the states that are written by the function.
 
@@ -316,6 +348,8 @@ class StatefulFunction(PrettyObject):
         Returns:
           The states that are written by the function.
         """
+        if cache_key is None:
+            cache_key = default_cache_key
         return self.get_state_trace(cache_key).get_write_states()
 
     def get_arg_cache_key(self, *args, **kwargs) -> Tuple:
@@ -323,10 +357,11 @@ class StatefulFunction(PrettyObject):
         Get the static arguments from the arguments.
 
         Args:
-          *args: The arguments to the function.
+            *args: The arguments to the function.
+            **kwargs: The keyword arguments to the function.
 
         Returns:
-          The static arguments.
+          The static arguments and keyword arguments as a tuple.
         """
         if self.cache_type == 'jit':
             static_args, dyn_args = [], []
@@ -336,11 +371,18 @@ class StatefulFunction(PrettyObject):
                 else:
                     dyn_args.append(arg)
             dyn_args = jax.tree.map(shaped_abstractify, jax.tree.leaves(dyn_args))
-            dyn_kwargs = jax.tree.map(shaped_abstractify, jax.tree.leaves(kwargs))
-            return tuple([tuple(static_args), tuple(dyn_args), tuple(dyn_kwargs)])
+            static_kwargs, dyn_kwargs = [], []
+            for k, v in kwargs.items():
+                if k in self.static_argnames:
+                    static_kwargs.append((k, v))
+                else:
+                    dyn_kwargs.append((k, jax.tree.map(shaped_abstractify, v)))
+            return tuple([tuple(static_args), tuple(dyn_args), tuple(static_kwargs), tuple(dyn_kwargs)])
         elif self.cache_type is None:
             num_arg = len(args)
-            return tuple(args[i] for i in self.static_argnums if i < num_arg)
+            static_args = tuple(args[i] for i in self.static_argnums if i < num_arg)
+            static_kwargs = tuple((k, v) for k, v in kwargs.items() if k in self.static_argnames)
+            return tuple([static_args, static_kwargs])
         else:
             raise ValueError(f"Invalid cache type: {self.cache_type}")
 
@@ -389,7 +431,7 @@ class StatefulFunction(PrettyObject):
         self._cached_state_trace.clear()
 
     def _wrapped_fun_to_eval(
-        self, cache_key, *args, return_only_write: bool = False, **kwargs,
+        self, cache_key, static_kwargs: dict, *args, return_only_write: bool = False, **dyn_kwargs,
     ) -> Tuple[Any, Tuple[State, ...]]:
         """
         Wrap the function and return the states that are read and written by the function and the output of the function.
@@ -405,7 +447,7 @@ class StatefulFunction(PrettyObject):
         state_trace = _init_state_trace_stack(self.name)
         self._cached_state_trace[cache_key] = state_trace
         with state_trace:
-            out = self.fun(*args, **kwargs)
+            out = self.fun(*args, **dyn_kwargs, **static_kwargs)
             state_values = (
                 state_trace.get_write_state_values(True)
                 if return_only_write else
@@ -430,8 +472,9 @@ class StatefulFunction(PrettyObject):
         the structure, shape, dtypes, and named shapes of the output of ``fun``.
 
         Args:
-          *args: The arguments to the function.
-          **kwargs: The keyword arguments to the function.
+            *args: The arguments to the function.
+            **kwargs: The keyword arguments to the function.
+            return_only_write: If True, only return the states that are written by the function.
         """
 
         # static args
@@ -440,17 +483,24 @@ class StatefulFunction(PrettyObject):
         if cache_key not in self._cached_state_trace:
             try:
                 # jaxpr
+                static_kwargs, dyn_kwargs = {}, {}
+                for k, v in kwargs.items():
+                    if k in self.static_argnames:
+                        static_kwargs[k] = v
+                    else:
+                        dyn_kwargs[k] = v
                 jaxpr, (out_shapes, state_shapes) = _make_jaxpr(
                     functools.partial(
                         self._wrapped_fun_to_eval,
                         cache_key,
+                        static_kwargs,
                         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)
+                )(*args, **dyn_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)
@@ -483,6 +533,7 @@ class StatefulFunction(PrettyObject):
         assert len(state_vals) == len(states), 'State length mismatch.'
 
         # parameters
+        kwargs = {k: v for k, v in kwargs.items() if k not in self.static_argnames}  # remove static kwargs
         args = tuple(args[i] for i in range(len(args)) if i not in self.static_argnums)
         args = jax.tree.flatten((args, kwargs, state_vals))[0]
 
@@ -519,12 +570,16 @@ class StatefulFunction(PrettyObject):
 def make_jaxpr(
     fun: Callable,
     static_argnums: Union[int, Iterable[int]] = (),
+    static_argnames: Union[str, Iterable[str]] = (),
     axis_env: Optional[Sequence[tuple[Hashable, int]]] = None,
     return_shape: bool = False,
     abstracted_axes: Optional[Any] = None,
     state_returns: Union[str, Tuple[str, ...]] = ('read', 'write')
-) -> Callable[..., (Tuple[ClosedJaxpr, Tuple[State, ...]] |
-                    Tuple[ClosedJaxpr, Tuple[State, ...], PyTree])]:
+) -> Callable[
+    ...,
+    (Tuple[ClosedJaxpr, Tuple[State, ...]] |
+     Tuple[ClosedJaxpr, Tuple[State, ...], PyTree])
+]:
     """
     Creates a function that produces its jaxpr given example args.
 
@@ -533,6 +588,7 @@ def make_jaxpr(
         arguments and return value should be arrays, scalars, or standard Python
         containers (tuple/list/dict) thereof.
       static_argnums: See the :py:func:`jax.jit` docstring.
+      static_argnames: See the :py:func:`jax.jit` docstring.
       axis_env: Optional, a sequence of pairs where the first element is an axis
         name and the second element is a positive integer representing the size of
         the mapped axis with that name. This parameter is useful when lowering
@@ -605,11 +661,11 @@ def make_jaxpr(
     stateful_fun = StatefulFunction(
         fun,
         static_argnums=static_argnums,
+        static_argnames=static_argnames,
         axis_env=axis_env,
         abstracted_axes=abstracted_axes,
         state_returns=state_returns,
         name='make_jaxpr'
-
     )
 
     @wraps(fun)

brainstate/compile/_make_jaxpr_test.py

@@ -88,7 +88,7 @@ class TestMakeJaxpr(unittest.TestCase):
         self.assertTrue(jnp.allclose(jaxpr_as_fun(jaxpr)(jnp.zeros(1), st1.value)[0],
                                      f3(jnp.zeros(1))))
 
-    def test_compar_jax_make_jaxpr2(self):
+    def test_compare_jax_make_jaxpr2(self):
         st1 = brainstate.State(jnp.ones(10))
 
         def fa(x):
@@ -108,7 +108,7 @@ class TestMakeJaxpr(unittest.TestCase):
         print(jaxpr)
         print(jaxpr_as_fun(jaxpr)(jnp.zeros(1)))
 
-    def test_compar_jax_make_jaxpr3(self):
+    def test_compare_jax_make_jaxpr3(self):
         def fa(x):
             return 1.
 
@@ -121,6 +121,17 @@ class TestMakeJaxpr(unittest.TestCase):
         print(jaxpr)
         # print(jaxpr_as_fun(jaxpr)(jnp.zeros(1)))
 
+    def test_static_argnames(self):
+        def func4(a, b):  # Arg is a pair
+            temp = a + jnp.sin(b) * 3.
+            c = brainstate.random.rand_like(a)
+            return jnp.sum(temp + c)
+
+        jaxpr, states = brainstate.compile.make_jaxpr(func4, static_argnames='b')(jnp.zeros(8), 1.)
+        print()
+        print(jaxpr)
+        print(states)
+
 
 def test_return_states():
     import jax.numpy

brainstate/graph/_graph_node.py

@@ -25,7 +25,7 @@ import numpy as np
 
 from brainstate._state import State, TreefyState
 from brainstate.typing import Key
-from brainstate.util._pretty_pytree import PrettyObject
+from brainstate.util.pretty_pytree import PrettyObject
 from ._graph_operation import register_graph_node_type
 
 __all__ = [

brainstate/graph/_graph_operation.py

@@ -30,10 +30,10 @@ from typing_extensions import TypeGuard, Unpack
 from brainstate._state import State, TreefyState
 from brainstate._utils import set_module_as
 from brainstate.typing import PathParts, Filter, Predicate, Key
-from brainstate.util._caller import ApplyCaller, CallableProxy, DelayedAccessor
-from brainstate.util._pretty_pytree import NestedDict, FlattedDict, PrettyDict
-from brainstate.util._pretty_repr import PrettyRepr, PrettyType, PrettyAttr, PrettyMapping, MappingReprMixin
-from brainstate.util._struct import FrozenDict
+from brainstate.util.caller import ApplyCaller, CallableProxy, DelayedAccessor
+from brainstate.util.pretty_pytree import NestedDict, FlattedDict, PrettyDict
+from brainstate.util.pretty_repr import PrettyRepr, PrettyType, PrettyAttr, PrettyMapping, MappingReprMixin
+from brainstate.util.struct import FrozenDict
 from brainstate.util.filter import to_predicate
 
 _max_int = np.iinfo(np.int32).max

brainstate/mixin.py

@@ -41,6 +41,14 @@ __all__ = [
 ]
 
 
+def hashable(x):
+    try:
+        hash(x)
+        return True
+    except TypeError:
+        return False
+
+
 class Mixin(object):
     """Base Mixin object.
 
@@ -67,6 +75,14 @@ class ParamDesc(Mixin):
 
 
 class HashableDict(dict):
+    def __init__(self, the_dict: dict):
+        out = dict()
+        for k, v in the_dict.items():
+            if not hashable(v):
+                v = str(v)  # convert to string if not hashable
+            out[k] = v
+        super().__init__(out)
+
     def __hash__(self):
         return hash(tuple(sorted(self.items())))
 
@@ -132,7 +148,6 @@ class AlignPost(Mixin):
         raise NotImplementedError
 
 
-
 class BindCondData(Mixin):
     """Bind temporary conductance data.
 
@@ -147,12 +162,26 @@ class BindCondData(Mixin):
         self._conductance = None
 
 
+def not_implemented(func):
+
+    def wrapper(*args, **kwargs):
+        raise NotImplementedError(f'{func.__name__} is not implemented.')
+
+    wrapper.not_implemented = True
+    return wrapper
+
+
+
 class UpdateReturn(Mixin):
+    @not_implemented
     def update_return(self) -> PyTree:
         """
         The update function return of the model.
 
-        It should be a pytree, with each element as a ``jax.ShapeDtypeStruct`` or ``jax.core.ShapedArray``.
+        This function requires no parameters and must return a PyTree.
+
+        It is usually used for delay initialization, for example, ``Dynamics.output_delay`` relies on this function to
+        initialize the output delay.
 
         """
         raise NotImplementedError(f'Must implement the "{self.update_return.__name__}()" function.')

brainstate/nn/__init__.py

@@ -33,12 +33,14 @@ from ._embedding import *
 from ._embedding import __all__ as embed_all
 from ._exp_euler import *
 from ._exp_euler import __all__ as exp_euler_all
-from ._fixedprob_mv import EventFixedProb, EventFixedNumConn
+from ._fixedprob import *
+from._fixedprob import __all__ as fixedprob_all
 from ._inputs import *
 from ._inputs import __all__ as inputs_all
 from ._linear import *
 from ._linear import __all__ as linear_all
-from ._linear_mv import EventLinear
+from ._linear_mv import *
+from ._linear_mv import __all__ as linear_mv_all
 from ._ltp import *
 from ._ltp import __all__ as ltp_all
 from ._module import *
@@ -69,9 +71,6 @@ from ._utils import __all__ as utils_all
 __all__ = (
     [
         'metrics',
-        'EventLinear',
-        'EventFixedProb',
-        'EventFixedNumConn',
     ]
     + collective_ops_all
     + common_all
@@ -87,6 +86,8 @@ __all__ = (
     + linear_all
     + normalizations_all
     + poolings_all
+    + fixedprob_all
+    + linear_mv_all
     + embed_all
     + dropout_all
     + elementwise_all
@@ -115,6 +116,8 @@ del (
     normalizations_all,
     poolings_all,
     embed_all,
+    fixedprob_all,
+    linear_mv_all,
     dropout_all,
     elementwise_all,
     dyn_neuron_all,

brainstate/nn/_delay.py

@@ -330,7 +330,14 @@ class Delay(Module):
         indices = (delay_idx,) + indices
 
         # the delay data
-        return jax.tree.map(lambda a: a[indices], self.history.value)
+        if self._unit is None:
+            return jax.tree.map(lambda a: a[indices], self.history.value)
+        else:
+            return jax.tree.map(
+                lambda hist, unit: u.maybe_decimal(hist[indices] * unit),
+                self.history.value,
+                self._unit
+            )
 
     def retrieve_at_time(self, delay_time, *indices) -> PyTree:
         """
@@ -393,6 +400,9 @@ class Delay(Module):
         """
         assert self.history is not None, 'The delay history is not initialized.'
 
+        if self.take_aware_unit and self._unit is None:
+            self._unit = jax.tree.map(lambda x: u.get_unit(x), current, is_leaf=u.math.is_quantity)
+
         # update the delay data at the rotation index
         if self.delay_method == _DELAY_ROTATE:
             i = environ.get(environ.I)
@@ -419,6 +429,8 @@ class Delay(Module):
             raise ValueError(f'Unknown updating method "{self.delay_method}"')
 
 
+
+
 class StateWithDelay(Delay):
     """
     A ``State`` type that defines the state in a differential equation.

brainstate/nn/_dropout.py

@@ -409,7 +409,8 @@ class DropoutFixed(ElementWiseBlock):
         self.out_size = in_size
 
     def init_state(self, batch_size=None, **kwargs):
-        self.mask = ShortTermState(init.param(partial(random.bernoulli, self.prob), self.in_size, batch_size))
+        if self.prob < 1.:
+            self.mask = ShortTermState(init.param(partial(random.bernoulli, self.prob), self.in_size, batch_size))
 
     def update(self, x):
         dtype = u.math.get_dtype(x)
@@ -418,8 +419,8 @@ class DropoutFixed(ElementWiseBlock):
             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.")
-            return jnp.where(self.mask.value,
-                             jnp.asarray(x / self.prob, dtype=dtype),
-                             jnp.asarray(0., dtype=dtype))
+            return u.math.where(self.mask.value,
+                                u.math.asarray(x / self.prob, dtype=dtype),
+                                u.math.asarray(0., dtype=dtype) * u.get_unit(x))
         else:
             return x