brainstate 0.0.1__py2.py3-none-any.whl → 0.0.1.post20240622__py2.py3-none-any.whl

brainstate/__init__.py

@@ -20,17 +20,16 @@ A ``State``-based Transformation System for Brain Dynamics Programming
 __version__ = "0.0.1"
 
 from . import environ
-from . import math
+from . import functional
+from . import init
 from . import mixin
 from . import nn
 from . import optim
 from . import random
+from . import surrogate
 from . import transform
 from . import typing
 from . import util
-from . import surrogate
-from . import functional
-from . import init
 from ._module import *
 from ._module import __all__ as _module_all
 from ._state import *
@@ -39,7 +38,7 @@ from ._state import __all__ as _state_all
 __all__ = (
     ['environ', 'share', 'nn', 'optim', 'random',
      'surrogate', 'functional', 'init',
-     'mixin', 'math', 'transform', 'util', 'typing'] +
+     'mixin', 'transform', 'util', 'typing'] +
     _module_all + _state_all
 )
 del _module_all, _state_all

brainstate/_module.py

@@ -46,7 +46,6 @@ For handling the delays:
 
 """
 
-import inspect
 import math
 import numbers
 from collections import namedtuple
@@ -58,20 +57,21 @@ import jax.numpy as jnp
 import numpy as np
 
 from . import environ
-from ._utils import set_module_as
 from ._state import State, StateDictManager, visible_state_dict
-from .util import unique_name, DictManager, get_unique_name, DotDict
-from .math import get_dtype
+from ._utils import set_module_as
 from .mixin import Mixin, Mode, DelayedInit, AllOfTypes, Batching, UpdateReturn
-from .transform._jit_error import jit_error
+from .transform import jit_error
+from .util import unique_name, DictManager, get_unique_name
 
 Shape = Union[int, Sequence[int]]
 PyTree = Any
 ArrayLike = jax.typing.ArrayLike
 
 delay_identifier = '_*_delay_of_'
-ROTATE_UPDATE = 'rotation'
-CONCAT_UPDATE = 'concat'
+_DELAY_ROTATE = 'rotation'
+_DELAY_CONCAT = 'concat'
+_INTERP_LINEAR = 'linear_interp'
+_INTERP_ROUND = 'round'
 
 StateLoadResult = namedtuple('StateLoadResult', ['missing_keys', 'unexpected_keys'])
 
@@ -92,7 +92,7 @@ __all__ = [
   'call_order',
 
   # state processing
-  'init_states', 'load_states', 'save_states', 'assign_state_values',
+  'init_states', 'reset_states', 'load_states', 'save_states', 'assign_state_values',
 ]
 
 
@@ -258,9 +258,8 @@ class Module(object):
     --------
 
     >>> import brainstate as bst
-    >>> import brainscale as nn  # noqa
     >>> x = bst.random.rand((10, 10))
-    >>> l = nn.Activation(jax.numpy.tanh)
+    >>> l = bst.nn.Activation(jax.numpy.tanh)
     >>> y = x >> l
     """
     return self.__call__(other)
@@ -271,6 +270,12 @@ class Module(object):
     """
     pass
 
+  def reset_state(self, *args, **kwargs):
+    """
+    State resetting function.
+    """
+    pass
+
   def save_state(self, **kwargs) -> Dict:
     """Save states as a dictionary. """
     return self.states(include_self=True, level=0, method='absolute')
@@ -396,8 +401,8 @@ class visible_module_list(list):
   retieved when using :py:func:`~.nodes()` function.
 
   >>> import brainstate as bst
-  >>> l = bst.visible_module_list([bp.dnn.Dense(1, 2),
-  >>>                     bp.dnn.LSTMCell(2, 3)])
+  >>> l = bst.visible_module_list([bst.nn.Linear(1, 2),
+  >>>                              bst.nn.LSTMCell(2, 3)])
   """
 
   __module__ = 'brainstate'
@@ -1033,14 +1038,14 @@ class Delay(ExtendedUpdateWithBA, DelayedInit):
          delay = length          data ]
     entries: optional, dict. The delay access entries.
     name: str. The delay name.
-    method: str. The method used for updating delay. Default None.
+    delay_method: str. The method used for updating delay. Default None.
     mode: Mode. The computing mode. Default None.
   """
 
   __module__ = 'brainstate'
 
-  non_hash_params = ('time', 'entries', 'name')
-  max_time: float
+  non_hashable_params = ('time', 'entries', 'name')
+  max_time: float  #
   max_length: int
   history: Optional[State]
 
@@ -1048,20 +1053,27 @@ class Delay(ExtendedUpdateWithBA, DelayedInit):
       self,
       target_info: PyTree,
       time: Optional[Union[int, float]] = None,  # delay time
-      init: Optional[Union[ArrayLike, Callable]] = None,  # delay data init
+      init: Optional[Union[ArrayLike, Callable]] = None,  # delay data before t0
       entries: Optional[Dict] = None,  # delay access entry
-      method: Optional[str] = ROTATE_UPDATE,  # delay method
+      delay_method: Optional[str] = _DELAY_ROTATE,  # delay method
+      interp_method: str = _INTERP_LINEAR,  # interpolation method
       # others
       name: Optional[str] = None,
       mode: Optional[Mode] = None,
   ):
 
     # target information
-    self.target_info = jax.tree.map(lambda a: jax.ShapeDtypeStruct(a.shape, get_dtype(a)), target_info)
+    self.target_info = jax.tree.map(lambda a: jax.ShapeDtypeStruct(a.shape, a.dtype), target_info)
 
     # delay method
-    assert method in [ROTATE_UPDATE, CONCAT_UPDATE]
-    self.method = method
+    assert delay_method in [_DELAY_ROTATE, _DELAY_CONCAT], (f'Un-supported delay method {delay_method}. '
+                                                            f'Only support {_DELAY_ROTATE} and {_DELAY_CONCAT}')
+    self.delay_method = delay_method
+
+    # interp method
+    assert interp_method in [_INTERP_LINEAR, _INTERP_ROUND], (f'Un-supported interpolation method {interp_method}. '
+                                                              f'we only support: {[_INTERP_LINEAR, _INTERP_ROUND]}')
+    self.interp_method = interp_method
 
     # delay length and time
     self.max_time, delay_length = _get_delay(time, None)
@@ -1071,7 +1083,8 @@ class Delay(ExtendedUpdateWithBA, DelayedInit):
 
     # delay data
     if init is not None:
-      assert isinstance(init, (numbers.Number, jax.Array, Callable))
+      if not isinstance(init, (numbers.Number, jax.Array, np.ndarray, Callable)):
+        raise TypeError(f'init should be Array, Callable, or None. But got {init}')
     self._init = init
     self._history = None
 
@@ -1085,7 +1098,11 @@ class Delay(ExtendedUpdateWithBA, DelayedInit):
 
   def __repr__(self):
     name = self.__class__.__name__
-    return f'{name}(delay_length={self.max_length}, target_info={self.target_info}, method="{self.method}")'
+    return (f'{name}('
+            f'delay_length={self.max_length}, '
+            f'target_info={self.target_info}, '
+            f'delay_method="{self.delay_method}", '
+            f'interp_method="{self.interp_method}")')
 
   @property
   def history(self):
@@ -1100,7 +1117,7 @@ class Delay(ExtendedUpdateWithBA, DelayedInit):
     if batch_size is not None:
       shape.insert(self.mode.batch_axis, batch_size)
     shape.insert(0, length)
-    if isinstance(self._init, (jax.Array, numbers.Number)):
+    if isinstance(self._init, (jax.Array, np.ndarray, numbers.Number)):
       data = jnp.broadcast_to(jnp.asarray(self._init, a.dtype), shape)
     elif callable(self._init):
       data = self._init(shape, dtype=a.dtype)
@@ -1115,13 +1132,20 @@ class Delay(ExtendedUpdateWithBA, DelayedInit):
     fun = partial(self._f_to_init, length=self.max_length, batch_size=batch_size)
     self.history = State(jax.tree.map(fun, self.target_info))
 
+  def reset_state(self, batch_size: int = None, **kwargs):
+    if batch_size is not None:
+      assert self.mode.has(Batching), 'The mode should have Batching behavior when batch_size is not None.'
+    fun = partial(self._f_to_init, length=self.max_length, batch_size=batch_size)
+    self.history.value = jax.tree.map(fun, self.target_info)
+
   def register_entry(
       self,
       entry: str,
       delay_time: Optional[Union[int, float]] = None,
       delay_step: Optional[int] = None,
   ) -> 'Delay':
-    """Register an entry to access the data.
+    """
+    Register an entry to access the delay data.
 
     Args:
       entry: str. The entry to access the delay data.
@@ -1151,7 +1175,8 @@ class Delay(ExtendedUpdateWithBA, DelayedInit):
     return self
 
   def at(self, entry: str, *indices) -> ArrayLike:
-    """Get the data at the given entry.
+    """
+    Get the data at the given entry.
 
     Args:
       entry: str. The entry to access the data.
@@ -1167,20 +1192,28 @@ class Delay(ExtendedUpdateWithBA, DelayedInit):
     delay_step = self._registered_entries[entry]
     if delay_step is None:
       delay_step = 0
-    return self.retrieve(delay_step, *indices)
+    return self.retrieve_at_step(delay_step, *indices)
 
-  def retrieve(self, delay_step, *indices):
-    """Retrieve the delay data according to the delay length.
+  def retrieve_at_step(self, delay_step, *indices) -> PyTree:
+    """
+    Retrieve the delay data at the given delay time step (the integer to indicate the time step).
 
     Parameters
     ----------
-    delay_step: int
-      The delay length used to retrieve the data.
+    delay_step: int_like
+      Retrieve the data at the given time step.
+    indices: tuple
+      The indices to slice the data.
+
+    Returns
+    -------
+    delay_data: The delay data at the given delay step.
+
     """
     assert self.history is not None, 'The delay history is not initialized.'
     assert delay_step is not None, 'The delay step should be given.'
 
-    if environ.get(environ.JIT_ERROR_CHECK, False):
+    if environ.get(environ.JIT_ERROR_CHECK, True):
       def _check_delay(delay_len):
         raise ValueError(f'The request delay length should be less than the '
                          f'maximum delay {self.max_length}. But we got {delay_len}')
@@ -1188,17 +1221,17 @@ class Delay(ExtendedUpdateWithBA, DelayedInit):
       jit_error(delay_step >= self.max_length, _check_delay, delay_step)
 
     # rotation method
-    if self.method == ROTATE_UPDATE:
-      i = environ.get(environ.I)
+    if self.delay_method == _DELAY_ROTATE:
+      i = environ.get(environ.I, desc='The time step index.')
       di = i - delay_step
       delay_idx = jnp.asarray(di % self.max_length, dtype=jnp.int32)
       delay_idx = jax.lax.stop_gradient(delay_idx)
 
-    elif self.method == CONCAT_UPDATE:
+    elif self.delay_method == _DELAY_CONCAT:
       delay_idx = delay_step
 
     else:
-      raise ValueError(f'Unknown updating method "{self.method}"')
+      raise ValueError(f'Unknown delay updating method "{self.delay_method}"')
 
     # the delay index
     if hasattr(delay_idx, 'dtype') and not jnp.issubdtype(delay_idx.dtype, jnp.integer):
@@ -1208,6 +1241,81 @@ class Delay(ExtendedUpdateWithBA, DelayedInit):
     # the delay data
     return jax.tree.map(lambda a: a[indices], self.history.value)
 
+  def retrieve_at_time(self, delay_time, *indices) -> PyTree:
+    """
+    Retrieve the delay data at the given delay time step (the integer to indicate the time step).
+
+    Parameters
+    ----------
+    delay_time: float
+      Retrieve the data at the given time.
+    indices: tuple
+      The indices to slice the data.
+
+    Returns
+    -------
+    delay_data: The delay data at the given delay step.
+
+    """
+    assert self.history is not None, 'The delay history is not initialized.'
+    assert delay_time is not None, 'The delay time should be given.'
+
+    current_time = environ.get(environ.T, desc='The current time.')
+    dt = environ.get_dt()
+
+    if environ.get(environ.JIT_ERROR_CHECK, True):
+      def _check_delay(args):
+        t_now, t_delay = args
+        raise ValueError(f'The request delay time should be within '
+                         f'[{t_now - self.max_time - dt}, {t_now}], '
+                         f'but we got {t_delay}')
+
+      jit_error(jnp.logical_or(delay_time > current_time,
+                               delay_time < current_time - self.max_time - dt),
+                _check_delay,
+                (current_time, delay_time))
+
+    diff = current_time - delay_time
+    float_time_step = diff / dt
+
+    if self.interp_method == _INTERP_LINEAR:  # "linear" interpolation
+      # def _interp(target):
+        # if len(indices) > 0:
+        #   raise NotImplementedError('The slicing indices are not supported in the linear interpolation.')
+        # if self.delay_method == _DELAY_ROTATE:
+        #   i = environ.get(environ.I, desc='The time step index.')
+        #   _interp_fun = partial(jnp.interp, period=self.max_length)
+        #   for dim in range(1, target.ndim, 1):
+        #     _interp_fun = jax.vmap(_interp_fun, in_axes=(None, None, dim), out_axes=dim - 1)
+        #   di = i - jnp.arange(self.max_length)
+        #   delay_idx = jnp.asarray(di % self.max_length, dtype=jnp.int32)
+        #   return _interp_fun(float_time_step, delay_idx, target)
+        #
+        # elif self.delay_method == _DELAY_CONCAT:
+        #   _interp_fun = partial(jnp.interp, period=self.max_length)
+        #   for dim in range(1, target.ndim, 1):
+        #     _interp_fun = jax.vmap(_interp_fun, in_axes=(None, None, dim), out_axes=dim - 1)
+        #   return _interp_fun(float_time_step, jnp.arange(self.max_length), target)
+        #
+        # else:
+        #   raise ValueError(f'Unknown delay updating method "{self.delay_method}"')
+      # return jax.tree.map(_interp, self.history.value)
+
+      data_at_t0 = self.retrieve_at_step(jnp.asarray(jnp.floor(float_time_step), dtype=jnp.int32), *indices)
+      data_at_t1 = self.retrieve_at_step(jnp.asarray(jnp.ceil(float_time_step), dtype=jnp.int32), *indices)
+      t_diff = float_time_step - jnp.floor(float_time_step)
+      return jax.tree.map(lambda a, b: a * (1 - t_diff) + b * t_diff, data_at_t0, data_at_t1)
+
+    elif self.interp_method == _INTERP_ROUND:  # "round" interpolation
+      return self.retrieve_at_step(
+        jnp.asarray(jnp.round(float_time_step), dtype=jnp.int32),
+        *indices
+      )
+
+    else:  # raise error
+      raise ValueError(f'Un-supported interpolation method {self.interp_method}, '
+                       f'we only support: {[_INTERP_LINEAR, _INTERP_ROUND]}')
+
   def update(self, current: PyTree) -> None:
     """
     Update delay variable with the new data.
@@ -1215,25 +1323,29 @@ class Delay(ExtendedUpdateWithBA, DelayedInit):
     assert self.history is not None, 'The delay history is not initialized.'
 
     # update the delay data at the rotation index
-    if self.method == ROTATE_UPDATE:
+    if self.delay_method == _DELAY_ROTATE:
       i = environ.get(environ.I)
       idx = jnp.asarray(i % self.max_length, dtype=environ.dutype())
       idx = jax.lax.stop_gradient(idx)
-      self.history.value = jax.tree.map(lambda hist, cur: hist.at[idx].set(cur),
-                                        self.history.value,
-                                        current)
+      self.history.value = jax.tree.map(
+        lambda hist, cur: hist.at[idx].set(cur),
+        self.history.value,
+        current
+      )
     # update the delay data at the first position
-    elif self.method == CONCAT_UPDATE:
+    elif self.delay_method == _DELAY_CONCAT:
       current = jax.tree.map(lambda a: jnp.expand_dims(a, 0), current)
       if self.max_length > 1:
-        self.history.value = jax.tree.map(lambda hist, cur: jnp.concatenate([cur, hist[:-1]], axis=0),
-                                          self.history.value,
-                                          current)
+        self.history.value = jax.tree.map(
+          lambda hist, cur: jnp.concatenate([cur, hist[:-1]], axis=0),
+          self.history.value,
+          current
+        )
       else:
         self.history.value = current
 
     else:
-      raise ValueError(f'Unknown updating method "{self.method}"')
+      raise ValueError(f'Unknown updating method "{self.delay_method}"')
 
 
 class _StateDelay(Delay):
@@ -1254,14 +1366,18 @@ class _StateDelay(Delay):
       time: Optional[Union[int, float]] = None,  # delay time
       init: Optional[Union[ArrayLike, Callable]] = None,  # delay data init
       entries: Optional[Dict] = None,  # delay access entry
-      method: Optional[str] = ROTATE_UPDATE,  # delay method
+      delay_method: Optional[str] = _DELAY_ROTATE,  # delay method
       # others
       name: Optional[str] = None,
       mode: Optional[Mode] = None,
   ):
     super().__init__(target_info=target.value,
-                     time=time, init=init, entries=entries,
-                     method=method, name=name, mode=mode)
+                     time=time,
+                     init=init,
+                     entries=entries,
+                     delay_method=delay_method,
+                     name=name,
+                     mode=mode)
     self.target = target
 
   def update(self, *args, **kwargs):
@@ -1344,7 +1460,7 @@ def call_order(level: int = 0):
 @set_module_as('brainstate')
 def init_states(target: Module, *args, **kwargs) -> Module:
   """
-  Reset states of all children nodes in the given target.
+  Initialize states of all children nodes in the given target.
 
   Args:
     target: The target Module.
@@ -1368,6 +1484,33 @@ def init_states(target: Module, *args, **kwargs) -> Module:
   return target
 
 
+@set_module_as('brainstate')
+def reset_states(target: Module, *args, **kwargs) -> Module:
+  """
+  Reset states of all children nodes in the given target.
+
+  Args:
+    target: The target Module.
+
+  Returns:
+    The target Module.
+  """
+  nodes_with_order = []
+
+  # reset node whose `init_state` has no `call_order`
+  for node in list(target.nodes().values()):
+    if not hasattr(node.reset_state, 'call_order'):
+      node.reset_state(*args, **kwargs)
+    else:
+      nodes_with_order.append(node)
+
+  # reset the node's states
+  for node in sorted(nodes_with_order, key=lambda x: x.reset_state.call_order):
+    node.reset_state(*args, **kwargs)
+
+  return target
+
+
 @set_module_as('brainstate')
 def load_states(target: Module, state_dict: Dict, **kwargs):
   """Copy parameters and buffers from :attr:`state_dict` into

brainstate/_module_test.py

@@ -16,14 +16,15 @@
 import unittest
 
 import jax.numpy as jnp
+import jaxlib.xla_extension
 
-import brainstate as bc
+import brainstate as bst
 
 
-class TestVarDelay(unittest.TestCase):
+class TestDelay(unittest.TestCase):
   def test_delay1(self):
-    a = bc.State(bc.random.random(10, 20))
-    delay = bc.Delay(a.value)
+    a = bst.State(bst.random.random(10, 20))
+    delay = bst.Delay(a.value)
     delay.register_entry('a', 1.)
     delay.register_entry('b', 2.)
     delay.register_entry('c', None)
@@ -31,10 +32,10 @@ class TestVarDelay(unittest.TestCase):
     delay.init_state()
     with self.assertRaises(KeyError):
       delay.register_entry('c', 10.)
-    bc.util.clear_buffer_memory()
+    bst.util.clear_buffer_memory()
 
   def test_rotation_delay(self):
-    rotation_delay = bc.Delay(jnp.ones((1,)))
+    rotation_delay = bst.Delay(jnp.ones((1,)))
     t0 = 0.
     t1, n1 = 1., 10
     t2, n2 = 2., 20
@@ -51,16 +52,16 @@ class TestVarDelay(unittest.TestCase):
     # print(rotation_delay.max_length)
 
     for i in range(100):
-      bc.environ.set(i=i)
+      bst.environ.set(i=i)
       rotation_delay(jnp.ones((1,)) * i)
       # print(i, rotation_delay.at('a'), rotation_delay.at('b'), rotation_delay.at('c2'), rotation_delay.at('c'))
       self.assertTrue(jnp.allclose(rotation_delay.at('a'), jnp.ones((1,)) * i))
       self.assertTrue(jnp.allclose(rotation_delay.at('b'), jnp.maximum(jnp.ones((1,)) * i - n1, 0.)))
       self.assertTrue(jnp.allclose(rotation_delay.at('c'), jnp.maximum(jnp.ones((1,)) * i - n2, 0.)))
-    bc.util.clear_buffer_memory()
+    bst.util.clear_buffer_memory()
 
   def test_concat_delay(self):
-    rotation_delay = bc.Delay(jnp.ones([1]), method='concat')
+    rotation_delay = bst.Delay(jnp.ones([1]), delay_method='concat')
     t0 = 0.
     t1, n1 = 1., 10
     t2, n2 = 2., 20
@@ -73,17 +74,91 @@ class TestVarDelay(unittest.TestCase):
 
     print()
     for i in range(100):
-      bc.environ.set(i=i)
+      bst.environ.set(i=i)
       rotation_delay(jnp.ones((1,)) * i)
       print(i, rotation_delay.at('a'), rotation_delay.at('b'), rotation_delay.at('c'))
       self.assertTrue(jnp.allclose(rotation_delay.at('a'), jnp.ones((1,)) * i))
       self.assertTrue(jnp.allclose(rotation_delay.at('b'), jnp.maximum(jnp.ones((1,)) * i - n1, 0.)))
       self.assertTrue(jnp.allclose(rotation_delay.at('c'), jnp.maximum(jnp.ones((1,)) * i - n2, 0.)))
-    bc.util.clear_buffer_memory()
+    bst.util.clear_buffer_memory()
+
+  def test_jit_erro(self):
+    rotation_delay = bst.Delay(jnp.ones([1]), time=2., delay_method='concat', interp_method='round')
+    rotation_delay.init_state()
+
+    with bst.environ.context(i=0, t=0):
+      rotation_delay.retrieve_at_time(-2.0)
+      with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
+        rotation_delay.retrieve_at_time(-2.1)
+      rotation_delay.retrieve_at_time(-2.01)
+      with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
+        rotation_delay.retrieve_at_time(-2.09)
+      with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
+        rotation_delay.retrieve_at_time(0.1)
+      with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
+        rotation_delay.retrieve_at_time(0.01)
+
+  def test_round_interp(self):
+    for shape in [(1,), (1, 1), (1, 1, 1)]:
+      for delay_method in ['rotation', 'concat']:
+        rotation_delay = bst.Delay(jnp.ones(shape), time=2., delay_method=delay_method, interp_method='round')
+        t0, n1 = 0.01, 0
+        t1, n1 = 1.04, 10
+        t2, n2 = 1.06, 11
+        rotation_delay.init_state()
+
+        @bst.transform.jit
+        def retrieve(td, i):
+          with bst.environ.context(i=i, t=i * bst.environ.get_dt()):
+            return rotation_delay.retrieve_at_time(td)
+
+        print()
+        for i in range(100):
+          t = i * bst.environ.get_dt()
+          with bst.environ.context(i=i, t=t):
+            rotation_delay(jnp.ones(shape) * i)
+            print(i,
+                  retrieve(t - t0, i),
+                  retrieve(t - t1, i),
+                  retrieve(t - t2, i))
+            self.assertTrue(jnp.allclose(retrieve(t - t0, i), jnp.ones(shape) * i))
+            self.assertTrue(jnp.allclose(retrieve(t - t1, i), jnp.maximum(jnp.ones(shape) * i - n1, 0.)))
+            self.assertTrue(jnp.allclose(retrieve(t - t2, i), jnp.maximum(jnp.ones(shape) * i - n2, 0.)))
+        bst.util.clear_buffer_memory()
+
+  def test_linear_interp(self):
+    for shape in [(1,), (1, 1), (1, 1, 1)]:
+      for delay_method in ['rotation', 'concat']:
+        print(shape, delay_method)
+
+        rotation_delay = bst.Delay(jnp.ones(shape), time=2., delay_method=delay_method, interp_method='linear_interp')
+        t0, n0 = 0.01, 0.1
+        t1, n1 = 1.04, 10.4
+        t2, n2 = 1.06, 10.6
+        rotation_delay.init_state()
+
+        @bst.transform.jit
+        def retrieve(td, i):
+          with bst.environ.context(i=i, t=i * bst.environ.get_dt()):
+            return rotation_delay.retrieve_at_time(td)
+
+        print()
+        for i in range(100):
+          t = i * bst.environ.get_dt()
+          with bst.environ.context(i=i, t=t):
+            rotation_delay(jnp.ones(shape) * i)
+            print(i,
+                  retrieve(t - t0, i),
+                  retrieve(t - t1, i),
+                  retrieve(t - t2, i))
+            self.assertTrue(jnp.allclose(retrieve(t - t0, i), jnp.maximum(jnp.ones(shape) * i - n0, 0.)))
+            self.assertTrue(jnp.allclose(retrieve(t - t1, i), jnp.maximum(jnp.ones(shape) * i - n1, 0.)))
+            self.assertTrue(jnp.allclose(retrieve(t - t2, i), jnp.maximum(jnp.ones(shape) * i - n2, 0.)))
+        bst.util.clear_buffer_memory()
 
   def test_rotation_and_concat_delay(self):
-    rotation_delay = bc.Delay(jnp.ones((1,)))
-    concat_delay = bc.Delay(jnp.ones([1]), method='concat')
+    rotation_delay = bst.Delay(jnp.ones((1,)))
+    concat_delay = bst.Delay(jnp.ones([1]), delay_method='concat')
     t0 = 0.
     t1, n1 = 1., 10
     t2, n2 = 2., 20
@@ -100,29 +175,29 @@ class TestVarDelay(unittest.TestCase):
 
     print()
     for i in range(100):
-      bc.environ.set(i=i)
+      bst.environ.set(i=i)
       new = jnp.ones((1,)) * i
       rotation_delay(new)
       concat_delay(new)
       self.assertTrue(jnp.allclose(rotation_delay.at('a'), concat_delay.at('a'), ))
       self.assertTrue(jnp.allclose(rotation_delay.at('b'), concat_delay.at('b'), ))
       self.assertTrue(jnp.allclose(rotation_delay.at('c'), concat_delay.at('c'), ))
-    bc.util.clear_buffer_memory()
+    bst.util.clear_buffer_memory()
 
 
 class TestModule(unittest.TestCase):
   def test_states(self):
-    class A(bc.Module):
+    class A(bst.Module):
       def __init__(self):
         super().__init__()
-        self.a = bc.State(bc.random.random(10, 20))
-        self.b = bc.State(bc.random.random(10, 20))
+        self.a = bst.State(bst.random.random(10, 20))
+        self.b = bst.State(bst.random.random(10, 20))
 
-    class B(bc.Module):
+    class B(bst.Module):
       def __init__(self):
         super().__init__()
         self.a = A()
-        self.b = bc.State(bc.random.random(10, 20))
+        self.b = bst.State(bst.random.random(10, 20))
 
     b = B()
     print()
@@ -130,4 +205,3 @@ class TestModule(unittest.TestCase):
     print(b.states())
     print(b.states(level=0))
     print(b.states(level=0))
-

brainstate/_state.py

@@ -59,6 +59,23 @@ _global_context_to_check_state_tree = [False]
 def check_state_value_tree() -> None:
   """
   The contex manager to check weather the tree structure of the state value keeps consistently.
+
+  Once a :py:class:`~.State` is created, the tree structure of the value is fixed. In default,
+  the tree structure of the value is not checked to avoid off the repeated evaluation.
+  If you want to check the tree structure of the value once the new value is assigned,
+  you can use this context manager.
+
+  Example::
+
+  ```python
+  state = brainstate.ShortTermState(jnp.zeros((2, 3)))
+  with check_state_value_tree():
+    state.value = jnp.zeros((2, 3))
+
+    # The following code will raise an error.
+    state.value = (jnp.zeros((2, 3)), jnp.zeros((2, 3)))
+  ```
+
   """
   try:
     _global_context_to_check_state_tree.append(True)