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

brainstate/compile/_make_jaxpr.py

@@ -352,6 +352,13 @@ class StatefulFunction(PrettyObject):
             cache_key = default_cache_key
         return self.get_state_trace(cache_key).get_write_states()
 
+    def _check_input(self, x):
+        if isinstance(x, State):
+            raise ValueError(
+                'Inputs for brainstate transformations cannot be an instance of State. '
+                f'But we got {x}'
+            )
+
     def get_arg_cache_key(self, *args, **kwargs) -> Tuple:
         """
         Get the static arguments from the arguments.
@@ -370,22 +377,35 @@ class StatefulFunction(PrettyObject):
                     static_args.append(arg)
                 else:
                     dyn_args.append(arg)
-            dyn_args = jax.tree.map(shaped_abstractify, jax.tree.leaves(dyn_args))
+            dyn_args = jax.tree.map(shaped_abstractify, dyn_args)
             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)])
+
+            static_args = make_hashable(tuple(static_args))
+            dyn_args = make_hashable(tuple(dyn_args))
+            static_kwargs = make_hashable(static_kwargs)
+            dyn_kwargs = make_hashable(dyn_kwargs)
+
+            cache_key = (static_args, dyn_args, static_kwargs, dyn_kwargs)
         elif self.cache_type is None:
             num_arg = len(args)
             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])
+
+            # Make everything hashable
+            static_args = make_hashable(static_args)
+            static_kwargs = make_hashable(static_kwargs)
+
+            cache_key = (static_args, static_kwargs)
         else:
             raise ValueError(f"Invalid cache type: {self.cache_type}")
 
+        return cache_key
+
     def compile_function_and_get_states(self, *args, **kwargs) -> Tuple[State, ...]:
         """
         Compile the function, and get the states that are read and written by this function.
@@ -480,6 +500,9 @@ class StatefulFunction(PrettyObject):
         # static args
         cache_key = self.get_arg_cache_key(*args, **kwargs)
 
+        # check input types
+        jax.tree.map(self._check_input, (args, kwargs), is_leaf=lambda x: isinstance(x, State))
+
         if cache_key not in self._cached_state_trace:
             try:
                 # jaxpr
@@ -637,15 +660,15 @@ def make_jaxpr(
     instead give a few examples.
 
     >>> import jax
-    >>> import brainstate as bst
+    >>> import brainstate as brainstate
     >>>
     >>> def f(x): return jax.numpy.sin(jax.numpy.cos(x))
     >>> print(f(3.0))
     -0.83602
-    >>> jaxpr, states = bst.compile.make_jaxpr(f)(3.0)
+    >>> jaxpr, states = brainstate.compile.make_jaxpr(f)(3.0)
     >>> jaxpr
     { lambda ; a:f32[]. let b:f32[] = cos a; c:f32[] = sin b in (c,) }
-    >>> jaxpr, states = bst.compile.make_jaxpr(jax.grad(f))(3.0)
+    >>> jaxpr, states = brainstate.compile.make_jaxpr(jax.grad(f))(3.0)
     >>> jaxpr
     { lambda ; a:f32[]. let
         b:f32[] = cos a
@@ -844,3 +867,33 @@ def _make_jaxpr(
     if hasattr(fun, "__name__"):
         make_jaxpr_f.__name__ = f"make_jaxpr({fun.__name__})"
     return make_jaxpr_f
+
+
+def make_hashable(obj):
+    """Convert a pytree into a hashable representation."""
+    if isinstance(obj, (list, tuple)):
+        return tuple(make_hashable(item) for item in obj)
+    elif isinstance(obj, dict):
+        return tuple(sorted((k, make_hashable(v)) for k, v in obj.items()))
+    elif isinstance(obj, set):
+        return frozenset(make_hashable(item) for item in obj)
+    elif hasattr(obj, '__dict__'):  # Handle custom objects
+        return (
+            obj.__class__.__name__,
+            tuple(
+                sorted(
+                    (k, make_hashable(v))
+                    for k, v in obj.__dict__.items()
+                    if not k.startswith('_')
+                )
+            )
+        )
+    else:
+        # # Use JAX's tree_util for any other pytree structures
+        # try:
+        #     leaves, treedef = jax.tree_util.tree_flatten(obj)
+        #     hashable_leaves = tuple(make_hashable(leaf) for leaf in leaves)
+        #     return (str(treedef), hashable_leaves)
+        # except:
+        #     # Assume obj is already hashable
+        return obj

brainstate/compile/_progress_bar.py

@@ -53,7 +53,7 @@ class ProgressBar(object):
 
     .. code-block:: python
 
-                a = bst.State(1.)
+                a = brainstate.State(1.)
                 def loop_fn(x):
                     a.value = x.value + 1.
                     return jnp.sum(x ** 2)
@@ -61,7 +61,7 @@ class ProgressBar(object):
                 pbar = ProgressBar(desc=("Running {i} iterations, loss = {loss}",
                                          lambda i_carray_y: {"i": i_carray_y["i"], "loss": i_carray_y["y"]}))
 
-                bst.compile.for_loop(loop_fn, xs, pbar=pbar)
+                brainstate.compile.for_loop(loop_fn, xs, pbar=pbar)
 
     In this example, ``"i"`` denotes the iteration number and ``"loss"`` is computed from the output,
     the ``"carry"`` is the dynamic state in the loop, for example ``a.value`` in this case.

brainstate/environ.py

@@ -76,9 +76,9 @@ def context(**kwargs):
 
     For instance::
 
-    >>> import brainstate as bst
-    >>> with bst.environ.context(dt=0.1) as env:
-    ...     dt = bst.environ.get('dt')
+    >>> import brainstate as brainstate
+    >>> with brainstate.environ.context(dt=0.1) as env:
+    ...     dt = brainstate.environ.get('dt')
     ...     print(env)
 
     """
@@ -424,10 +424,10 @@ def dftype() -> DTypeLike:
 
     For example, if the precision is set to 32, the default floating data type is ``np.float32``.
 
-    >>> import brainstate as bst
+    >>> import brainstate as brainstate
     >>> import numpy as np
-    >>> with bst.environ.context(precision=32):
-    ...    a = np.zeros(1, dtype=bst.environ.dftype())
+    >>> with brainstate.environ.context(precision=32):
+    ...    a = np.zeros(1, dtype=brainstate.environ.dftype())
     >>> print(a.dtype)
 
     Returns
@@ -448,10 +448,10 @@ def ditype() -> DTypeLike:
 
     For example, if the precision is set to 32, the default integer data type is ``np.int32``.
 
-    >>> import brainstate as bst
+    >>> import brainstate as brainstate
     >>> import numpy as np
-    >>> with bst.environ.context(precision=32):
-    ...    a = np.zeros(1, dtype=bst.environ.ditype())
+    >>> with brainstate.environ.context(precision=32):
+    ...    a = np.zeros(1, dtype=brainstate.environ.ditype())
     >>> print(a.dtype)
     int32
 
@@ -474,10 +474,10 @@ def dutype() -> DTypeLike:
 
     For example, if the precision is set to 32, the default unsigned integer data type is ``np.uint32``.
 
-    >>> import brainstate as bst
+    >>> import brainstate as brainstate
     >>> import numpy as np
-    >>> with bst.environ.context(precision=32):
-    ...   a = np.zeros(1, dtype=bst.environ.dutype())
+    >>> with brainstate.environ.context(precision=32):
+    ...   a = np.zeros(1, dtype=brainstate.environ.dutype())
     >>> print(a.dtype)
     uint32
 
@@ -499,10 +499,10 @@ def dctype() -> DTypeLike:
 
     For example, if the precision is set to 32, the default complex data type is ``np.complex64``.
 
-    >>> import brainstate as bst
+    >>> import brainstate as brainstate
     >>> import numpy as np
-    >>> with bst.environ.context(precision=32):
-    ...    a = np.zeros(1, dtype=bst.environ.dctype())
+    >>> with brainstate.environ.context(precision=32):
+    ...    a = np.zeros(1, dtype=brainstate.environ.dctype())
     >>> print(a.dtype)
     complex64
 
@@ -529,19 +529,19 @@ def register_default_behavior(key: str, behavior: Callable, replace_if_exist: bo
 
     For example, you can register a default behavior for the key 'dt' by::
 
-    >>> import brainstate as bst
+    >>> import brainstate as brainstate
     >>> def dt_behavior(dt):
     ...     print(f'Set the default dt to {dt}.')
     ...
-    >>> bst.environ.register_default_behavior('dt', dt_behavior)
+    >>> brainstate.environ.register_default_behavior('dt', dt_behavior)
 
     Then, when you set the default dt by `brainstate.environ.set(dt=0.1)`, the behavior
     `dt_behavior` will be called with
     `dt_behavior(0.1)`.
 
-    >>> bst.environ.set(dt=0.1)
+    >>> brainstate.environ.set(dt=0.1)
     Set the default dt to 0.1.
-    >>> with bst.environ.context(dt=0.2):
+    >>> with brainstate.environ.context(dt=0.2):
     ...     pass
     Set the default dt to 0.2.
     Set the default dt to 0.1.

brainstate/functional/_activations_test.py

@@ -70,39 +70,39 @@ class NNFunctionsTest(jtu.JaxTestCase):
             check_dtypes=False)
 
     #   def testSquareplusGrad(self):
-    #     check_grads(bst.functional.squareplus, (1e-8,), order=4,
+    #     check_grads(brainstate.functional.squareplus, (1e-8,), order=4,
     #                 )
 
     #   def testSquareplusGradZero(self):
-    #     check_grads(bst.functional.squareplus, (0.,), order=1,
+    #     check_grads(brainstate.functional.squareplus, (0.,), order=1,
     #                 )
 
     #   def testSquareplusGradNegInf(self):
-    #     check_grads(bst.functional.squareplus, (-float('inf'),), order=1,
+    #     check_grads(brainstate.functional.squareplus, (-float('inf'),), order=1,
     #                 )
 
     #   def testSquareplusGradNan(self):
-    #     check_grads(bst.functional.squareplus, (float('nan'),), order=1,
+    #     check_grads(brainstate.functional.squareplus, (float('nan'),), order=1,
     #                 )
 
     #   @parameterized.parameters([float] + jtu.dtypes.floating)
     #   def testSquareplusZero(self, dtype):
-    #     self.assertEqual(dtype(1), bst.functional.squareplus(dtype(0), dtype(4)))
+    #     self.assertEqual(dtype(1), brainstate.functional.squareplus(dtype(0), dtype(4)))
     #
     # def testMishGrad(self):
-    #   check_grads(bst.functional.mish, (1e-8,), order=4,
+    #   check_grads(brainstate.functional.mish, (1e-8,), order=4,
     #               )
     #
     # def testMishGradZero(self):
-    #   check_grads(bst.functional.mish, (0.,), order=1,
+    #   check_grads(brainstate.functional.mish, (0.,), order=1,
     #               )
     #
     # def testMishGradNegInf(self):
-    #   check_grads(bst.functional.mish, (-float('inf'),), order=1,
+    #   check_grads(brainstate.functional.mish, (-float('inf'),), order=1,
     #               )
     #
     # def testMishGradNan(self):
-    #   check_grads(bst.functional.mish, (float('nan'),), order=1,
+    #   check_grads(brainstate.functional.mish, (float('nan'),), order=1,
     #               )
 
     @parameterized.parameters([float] + jtu.dtypes.floating)
@@ -137,7 +137,7 @@ class NNFunctionsTest(jtu.JaxTestCase):
         self.assertAllClose(brainstate.functional.sparse_sigmoid(0.), .5, check_dtypes=False)
 
     #   def testSquareplusValue(self):
-    #     val = bst.functional.squareplus(1e3)
+    #     val = brainstate.functional.squareplus(1e3)
     #     self.assertAllClose(val, 1e3, check_dtypes=False, atol=1e-3)
 
     def testMishValue(self):
@@ -177,7 +177,7 @@ class NNFunctionsTest(jtu.JaxTestCase):
          brainstate.functional.softplus,
          brainstate.functional.sparse_plus,
          brainstate.functional.sigmoid,
-         #  bst.functional.squareplus,
+         #  brainstate.functional.squareplus,
          brainstate.functional.mish)))
     def testDtypeMatchesInput(self, dtype, fn):
         x = jnp.zeros((), dtype=dtype)
@@ -306,7 +306,7 @@ class NNFunctionsTest(jtu.JaxTestCase):
 
     def testCustomJVPLeak2(self):
         # https://github.com/google/jax/issues/8171
-        # The above test uses jax.bst.functional.sigmoid, as in the original #8171, but that
+        # The above test uses jax.brainstate.functional.sigmoid, as in the original #8171, but that
         # function no longer actually has a custom_jvp! So we inline the old def.
 
         @jax.custom_jvp

brainstate/graph/_graph_operation.py

@@ -473,8 +473,8 @@ def flatten(
 
     Example::
 
-        >>> import brainstate as bst
-        >>> node = bst.graph.Node()
+        >>> import brainstate as brainstate
+        >>> node = brainstate.graph.Node()
         >>> graph_def, state_mapping = flatten(node)
         >>> print(graph_def)
         >>> print(state_mapping)
@@ -709,15 +709,15 @@ def unflatten(
 
     Example::
 
-    >>> import brainstate as bst
-    >>> class MyNode(bst.graph.Node):
+    >>> import brainstate as brainstate
+    >>> class MyNode(brainstate.graph.Node):
     ...   def __init__(self):
-    ...      self.a = bst.nn.Linear(2, 3)
-    ...      self.b = bst.nn.Linear(3, 4)
-    ...      self.c = [bst.nn.Linear(4, 5), bst.nn.Linear(5, 6)]
-    ...      self.d = {'x': bst.nn.Linear(6, 7), 'y': bst.nn.Linear(7, 8)}
+    ...      self.a = brainstate.nn.Linear(2, 3)
+    ...      self.b = brainstate.nn.Linear(3, 4)
+    ...      self.c = [brainstate.nn.Linear(4, 5), brainstate.nn.Linear(5, 6)]
+    ...      self.d = {'x': brainstate.nn.Linear(6, 7), 'y': brainstate.nn.Linear(7, 8)}
     ...
-    >>> graphdef, statetree = bst.graph.flatten(MyNode())
+    >>> graphdef, statetree = brainstate.graph.flatten(MyNode())
     >>> statetree
     NestedDict({
       'a': {
@@ -764,7 +764,7 @@ def unflatten(
         }
       }
     })
-    >>> node = bst.graph.unflatten(graphdef, statetree)
+    >>> node = brainstate.graph.unflatten(graphdef, statetree)
     >>> node
     MyNode(
       a=Linear(
@@ -942,21 +942,21 @@ def pop_states(
 
     Example usage::
 
-      >>> import brainstate as bst
+      >>> import brainstate as brainstate
       >>> import jax.numpy as jnp
 
-      >>> class Model(bst.nn.Module):
+      >>> class Model(brainstate.nn.Module):
       ...   def __init__(self):
       ...     super().__init__()
-      ...     self.a = bst.nn.Linear(2, 3)
-      ...     self.b = bst.nn.LIF([10, 2])
+      ...     self.a = brainstate.nn.Linear(2, 3)
+      ...     self.b = brainstate.nn.LIF([10, 2])
 
       >>> model = Model()
-      >>> with bst.catch_new_states('new'):
-      ...    bst.nn.init_all_states(model)
+      >>> with brainstate.catch_new_states('new'):
+      ...    brainstate.nn.init_all_states(model)
 
       >>> assert len(model.states()) == 2
-      >>> model_states = bst.graph.pop_states(model, 'new')
+      >>> model_states = brainstate.graph.pop_states(model, 'new')
       >>> model_states
       NestedDict({
         'b': {
@@ -1046,16 +1046,16 @@ def treefy_split(
 
     Example usage::
 
-      >>> from joblib.testing import param    >>> import brainstate as bst
+      >>> from joblib.testing import param    >>> import brainstate as brainstate
       >>> import jax, jax.numpy as jnp
       ...
-      >>> class Foo(bst.graph.Node):
+      >>> class Foo(brainstate.graph.Node):
       ...   def __init__(self):
-      ...     self.a = bst.nn.BatchNorm1d([10, 2])
-      ...     self.b = bst.nn.Linear(2, 3)
+      ...     self.a = brainstate.nn.BatchNorm1d([10, 2])
+      ...     self.b = brainstate.nn.Linear(2, 3)
       ...
       >>> node = Foo()
-      >>> graphdef, params, others = bst.graph.treefy_split(node, bst.ParamState, ...)
+      >>> graphdef, params, others = brainstate.graph.treefy_split(node, brainstate.ParamState, ...)
       ...
       >>> params
       NestedDict({
@@ -1119,21 +1119,21 @@ def treefy_merge(
 
     Example usage::
 
-      >>> import brainstate as bst
+      >>> import brainstate as brainstate
       >>> import jax, jax.numpy as jnp
       ...
-      >>> class Foo(bst.graph.Node):
+      >>> class Foo(brainstate.graph.Node):
       ...   def __init__(self):
-      ...     self.a = bst.nn.BatchNorm1d([10, 2])
-      ...     self.b = bst.nn.Linear(2, 3)
+      ...     self.a = brainstate.nn.BatchNorm1d([10, 2])
+      ...     self.b = brainstate.nn.Linear(2, 3)
       ...
       >>> node = Foo()
-      >>> graphdef, params, others = bst.graph.treefy_split(node, bst.ParamState, ...)
+      >>> graphdef, params, others = brainstate.graph.treefy_split(node, brainstate.ParamState, ...)
       ...
-      >>> new_node = bst.graph.treefy_merge(graphdef, params, others)
+      >>> new_node = brainstate.graph.treefy_merge(graphdef, params, others)
       >>> assert isinstance(new_node, Foo)
-      >>> assert isinstance(new_node.b, bst.nn.BatchNorm1d)
-      >>> assert isinstance(new_node.a, bst.nn.Linear)
+      >>> assert isinstance(new_node.b, brainstate.nn.BatchNorm1d)
+      >>> assert isinstance(new_node.a, brainstate.nn.Linear)
 
     :func:`split` and :func:`merge` are primarily used to interact directly with JAX
     transformations, see
@@ -1302,22 +1302,22 @@ def treefy_states(
 
     Example usage::
 
-      >>> import brainstate as bst
-      >>> class Model(bst.nn.Module):
+      >>> import brainstate as brainstate
+      >>> class Model(brainstate.nn.Module):
       ...   def __init__(self):
       ...     super().__init__()
-      ...     self.l1 = bst.nn.Linear(2, 3)
-      ...     self.l2 = bst.nn.Linear(3, 4)
+      ...     self.l1 = brainstate.nn.Linear(2, 3)
+      ...     self.l2 = brainstate.nn.Linear(3, 4)
       ...   def __call__(self, x):
       ...     return self.l2(self.l1(x))
 
       >>> model = Model()
       >>> # get the learnable parameters from the batch norm and linear layer
-      >>> params = bst.graph.treefy_states(model, bst.ParamState)
+      >>> params = brainstate.graph.treefy_states(model, brainstate.ParamState)
       >>> # get them separately
-      >>> params, others = bst.graph.treefy_states(model, bst.ParamState, bst.ShortTermState)
+      >>> params, others = brainstate.graph.treefy_states(model, brainstate.ParamState, brainstate.ShortTermState)
       >>> # get them together
-      >>> states = bst.graph.treefy_states(model)
+      >>> states = brainstate.graph.treefy_states(model)
 
     Args:
       node: A graph node object.
@@ -1403,11 +1403,11 @@ def graphdef(node: Any, /) -> GraphDef[Any]:
 
     Example usage::
 
-      >>> import brainstate as bst
+      >>> import brainstate as brainstate
 
-      >>> model = bst.nn.Linear(2, 3)
-      >>> graphdef, _ = bst.graph.treefy_split(model)
-      >>> assert graphdef == bst.graph.graphdef(model)
+      >>> model = brainstate.nn.Linear(2, 3)
+      >>> graphdef, _ = brainstate.graph.treefy_split(model)
+      >>> assert graphdef == brainstate.graph.graphdef(model)
 
     Args:
       node: A graph node object.
@@ -1426,8 +1426,8 @@ def clone(node: Node) -> Node:
 
     Example usage::
 
-      >>> import brainstate as bst
-      >>> model = bst.nn.Linear(2, 3)
+      >>> import brainstate as brainstate
+      >>> model = brainstate.nn.Linear(2, 3)
       >>> cloned_model = clone(model)
       >>> model.weight.value['bias'] += 1
       >>> assert (model.weight.value['bias'] != cloned_model.weight.value['bias']).all()
@@ -1456,15 +1456,15 @@ def call(
 
     Example::
 
-      >>> import brainstate as bst
+      >>> import brainstate as brainstate
       >>> import jax
       >>> import jax.numpy as jnp
       ...
-      >>> class StatefulLinear(bst.graph.Node):
+      >>> class StatefulLinear(brainstate.graph.Node):
       ...   def __init__(self, din, dout):
-      ...     self.w = bst.ParamState(bst.random.rand(din, dout))
-      ...     self.b = bst.ParamState(jnp.zeros((dout,)))
-      ...     self.count = bst.State(jnp.array(0, dtype=jnp.uint32))
+      ...     self.w = brainstate.ParamState(brainstate.random.rand(din, dout))
+      ...     self.b = brainstate.ParamState(jnp.zeros((dout,)))
+      ...     self.count = brainstate.State(jnp.array(0, dtype=jnp.uint32))
       ...
       ...   def increment(self):
       ...     self.count.value += 1
@@ -1474,18 +1474,18 @@ def call(
       ...     return x @ self.w.value + self.b.value
       ...
       >>> linear = StatefulLinear(3, 2)
-      >>> linear_state = bst.graph.treefy_split(linear)
+      >>> linear_state = brainstate.graph.treefy_split(linear)
       ...
       >>> @jax.jit
       ... def forward(x, linear_state):
-      ...   y, linear_state = bst.graph.call(linear_state)(x)
+      ...   y, linear_state = brainstate.graph.call(linear_state)(x)
       ...   return y, linear_state
       ...
       >>> x = jnp.ones((1, 3))
       >>> y, linear_state = forward(x, linear_state)
       >>> y, linear_state = forward(x, linear_state)
       ...
-      >>> linear = bst.graph.treefy_merge(*linear_state)
+      >>> linear = brainstate.graph.treefy_merge(*linear_state)
       >>> linear.count.value
       Array(2, dtype=uint32)
 
@@ -1494,11 +1494,11 @@ def call(
     is used to call the ``increment`` method of the ``StatefulLinear`` module
     at the ``b`` key of a ``nodes`` dictionary.
 
-      >>> class StatefulLinear(bst.graph.Node):
+      >>> class StatefulLinear(brainstate.graph.Node):
       ...   def __init__(self, din, dout):
-      ...     self.w = bst.ParamState(bst.random.rand(din, dout))
-      ...     self.b = bst.ParamState(jnp.zeros((dout,)))
-      ...     self.count = bst.State(jnp.array(0, dtype=jnp.uint32))
+      ...     self.w = brainstate.ParamState(brainstate.random.rand(din, dout))
+      ...     self.b = brainstate.ParamState(jnp.zeros((dout,)))
+      ...     self.count = brainstate.State(jnp.array(0, dtype=jnp.uint32))
       ...
       ...   def increment(self):
       ...     self.count.value += 1
@@ -1514,7 +1514,7 @@ def call(
       ...
       >>> node_state = treefy_split(nodes)
       >>> # use attribute access
-      >>> _, node_state = bst.graph.call(node_state)['b'].increment()
+      >>> _, node_state = brainstate.graph.call(node_state)['b'].increment()
       ...
       >>> nodes = treefy_merge(*node_state)
       >>> nodes['a'].count.value
@@ -1544,19 +1544,19 @@ def iter_leaf(
     root. Repeated nodes are visited only once. Leaves include static values.
 
     Example::
-      >>> import brainstate as bst
+      >>> import brainstate as brainstate
       >>> import jax.numpy as jnp
       ...
-      >>> class Linear(bst.nn.Module):
+      >>> class Linear(brainstate.nn.Module):
       ...   def __init__(self, din, dout):
       ...     super().__init__()
-      ...     self.weight = bst.ParamState(bst.random.randn(din, dout))
-      ...     self.bias = bst.ParamState(bst.random.randn(dout))
+      ...     self.weight = brainstate.ParamState(brainstate.random.randn(din, dout))
+      ...     self.bias = brainstate.ParamState(brainstate.random.randn(dout))
       ...     self.a = 1
       ...
       >>> module = Linear(3, 4)
       ...
-      >>> for path, value in bst.graph.iter_leaf([module, module]):
+      >>> for path, value in brainstate.graph.iter_leaf([module, module]):
       ...   print(path, type(value).__name__)
       ...
       (0, 'a') int
@@ -1616,21 +1616,21 @@ def iter_node(
     root. Repeated nodes are visited only once. Leaves include static values.
 
     Example::
-      >>> import brainstate as bst
+      >>> import brainstate as brainstate
       >>> import jax.numpy as jnp
       ...
-      >>> class Model(bst.nn.Module):
+      >>> class Model(brainstate.nn.Module):
       ...   def __init__(self):
       ...     super().__init__()
-      ...     self.a = bst.nn.Linear(1, 2)
-      ...     self.b = bst.nn.Linear(2, 3)
-      ...     self.c = [bst.nn.Linear(3, 4), bst.nn.Linear(4, 5)]
-      ...     self.d = {'x': bst.nn.Linear(5, 6), 'y': bst.nn.Linear(6, 7)}
-      ...     self.b.a = bst.nn.LIF(2)
+      ...     self.a = brainstate.nn.Linear(1, 2)
+      ...     self.b = brainstate.nn.Linear(2, 3)
+      ...     self.c = [brainstate.nn.Linear(3, 4), brainstate.nn.Linear(4, 5)]
+      ...     self.d = {'x': brainstate.nn.Linear(5, 6), 'y': brainstate.nn.Linear(6, 7)}
+      ...     self.b.a = brainstate.nn.LIF(2)
       ...
       >>> model = Model()
       ...
-      >>> for path, node in bst.graph.iter_node([model, model]):
+      >>> for path, node in brainstate.graph.iter_node([model, model]):
       ...    print(path, node.__class__.__name__)
       ...
       (0, 'a') Linear

brainstate/graph/_graph_operation_test.py

@@ -443,7 +443,7 @@ class TestGraphOperation(unittest.TestCase):
         graphdef, statetree = brainstate.graph.flatten(MyNode())
         # print(graphdef)
         print(statetree)
-        # print(bst.graph.unflatten(graphdef, statetree))
+        # print(brainstate.graph.unflatten(graphdef, statetree))
 
     def test_split(self):
         class Foo(brainstate.graph.Node):
@@ -530,7 +530,7 @@ class TestGraphOperation(unittest.TestCase):
         print(graph_def)
         print(treefy_states)
 
-        # states = bst.graph.states(model)
+        # states = brainstate.graph.states(model)
         # print(states)
         # nest_states = states.to_nest()
         # print(nest_states)

brainstate/mixin.py

@@ -27,7 +27,6 @@ __all__ = [
     'ParamDescriber',
     'AlignPost',
     'BindCondData',
-    'UpdateReturn',
 
     # types
     'JointTypes',
@@ -171,22 +170,6 @@ def not_implemented(func):
     return wrapper
 
 
-
-class UpdateReturn(Mixin):
-    @not_implemented
-    def update_return(self) -> PyTree:
-        """
-        The update function return of the model.
-
-        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.')
-
-
 class _MetaUnionType(type):
     def __new__(cls, name, bases, dct):
         if isinstance(bases, type):

brainstate/nn/_collective_ops.py

@@ -56,10 +56,10 @@ def call_order(level: int = 0, check_order_boundary: bool = True):
 
     The lower the level, the earlier the function is called.
 
-    >>> import brainstate as bst
-    >>> bst.nn.call_order(0)
-    >>> bst.nn.call_order(-1)
-    >>> bst.nn.call_order(-2)
+    >>> import brainstate as brainstate
+    >>> brainstate.nn.call_order(0)
+    >>> brainstate.nn.call_order(-1)
+    >>> brainstate.nn.call_order(-2)
 
     Parameters
     ----------