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

brainstate/_module.py

@@ -46,7 +46,6 @@ For handling the delays:
 
 """
 
-import inspect
 import math
 import numbers
 from collections import namedtuple
@@ -58,12 +57,12 @@ 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 ._utils import set_module_as
 from .math import get_dtype
 from .mixin import Mixin, Mode, DelayedInit, AllOfTypes, Batching, UpdateReturn
 from .transform._jit_error import jit_error
+from .util import unique_name, DictManager, get_unique_name
 
 Shape = Union[int, Sequence[int]]
 PyTree = Any
@@ -92,7 +91,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',
 ]
 
 
@@ -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')
@@ -1115,6 +1120,12 @@ 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,
@@ -1344,7 +1355,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 +1379,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/_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)

brainstate/environ.py

@@ -18,7 +18,8 @@ from .util import MemScaling, IdMemScaling
 __all__ = [
   'set', 'context', 'get', 'all',
   'set_host_device_count', 'set_platform',
-  'get_host_device_count', 'get_platform', 'get_dt', 'get_mode', 'get_mem_scaling', 'get_precision',
+  'get_host_device_count', 'get_platform',
+  'get_dt', 'get_mode', 'get_mem_scaling', 'get_precision',
   'tolerance',
   'dftype', 'ditype', 'dutype', 'dctype',
 ]

brainstate/functional/__init__.py

@@ -20,6 +20,7 @@ from ._normalization import *
 from ._normalization import __all__ as __others_all__
 from ._spikes import *
 from ._spikes import __all__ as __spikes_all__
+from ._others import *
+from ._others import __all__ as __others_all__
 
-__all__ = __spikes_all__ + __others_all__ + __activations_all__
-
+__all__ = __spikes_all__ + __others_all__ + __activations_all__ + __others_all__

brainstate/functional/_activations.py

@@ -27,7 +27,7 @@ import jax.numpy as jnp
 from jax.scipy.special import logsumexp
 from jax.typing import ArrayLike
 
-from brainstate import math, random
+from .. import math, random
 
 __all__ = [
   "tanh",

brainstate/functional/_normalization.py

@@ -20,11 +20,14 @@ from typing import Optional
 import jax
 import jax.numpy as jnp
 
+from .._utils import set_module_as
+
 __all__ = [
   'weight_standardization',
 ]
 
 
+@set_module_as('brainstate.functional')
 def weight_standardization(
     w: jax.typing.ArrayLike,
     eps: float = 1e-4,

brainstate/functional/_others.py

@@ -0,0 +1,49 @@
+# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from __future__ import annotations
+
+from functools import partial
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+
+PyTree = Any
+
+__all__ = [
+  'clip_grad_norm',
+]
+
+
+def clip_grad_norm(
+    grad: PyTree,
+    max_norm: float | jax.Array,
+    norm_type: int | str | None = None
+):
+  """
+  Clips gradient norm of an iterable of parameters.
+
+  The norm is computed over all gradients together, as if they were
+  concatenated into a single vector. Gradients are modified in-place.
+
+  Args:
+      grad (PyTree): an iterable of Tensors or a single Tensor that will have gradients normalized
+      max_norm (float): max norm of the gradients.
+      norm_type (int, str, None): type of the used p-norm. Can be ``'inf'`` for infinity norm.
+  """
+  norm_fn = partial(jnp.linalg.norm, ord=norm_type)
+  norm = norm_fn(jnp.asarray(jax.tree.leaves(jax.tree.map(norm_fn, grad))))
+  return jax.tree.map(lambda x: jnp.where(norm < max_norm, x, x * max_norm / (norm + 1e-6)), grad)

brainstate/nn/__init__.py

@@ -21,6 +21,8 @@ from ._dynamics import *
 from ._dynamics import __all__ as dynamics_all
 from ._elementwise import *
 from ._elementwise import __all__ as elementwise_all
+from ._embedding import *
+from ._embedding import __all__ as embed_all
 from ._misc import *
 from ._misc import __all__ as _misc_all
 from ._normalizations import *
@@ -43,6 +45,7 @@ __all__ = (
     connections_all +
     dynamics_all +
     elementwise_all +
+    embed_all +
     normalizations_all +
     others_all +
     poolings_all +
@@ -58,6 +61,7 @@ del (
   connections_all,
   dynamics_all,
   elementwise_all,
+  embed_all,
   normalizations_all,
   others_all,
   poolings_all,

brainstate/nn/_base.py

@@ -55,22 +55,24 @@ class ExplicitInOutSize(Mixin):
 
   @property
   def in_size(self) -> Tuple[int, ...]:
-    if self._in_size is None:
-      raise ValueError(f"The input shape is not set in this node: {self} ")
     return self._in_size
 
   @in_size.setter
-  def in_size(self, in_size: Sequence[int]):
+  def in_size(self, in_size: Sequence[int] | int):
+    if isinstance(in_size, int):
+      in_size = (in_size,)
+    assert isinstance(in_size, (tuple, list)), f"Invalid type of in_size: {type(in_size)}"
     self._in_size = tuple(in_size)
 
   @property
   def out_size(self) -> Tuple[int, ...]:
-    if self._out_size is None:
-      raise ValueError(f"The output shape is not set in this node: {self}")
     return self._out_size
 
   @out_size.setter
-  def out_size(self, out_size: Sequence[int]):
+  def out_size(self, out_size: Sequence[int] | int):
+    if isinstance(out_size, int):
+      out_size = (out_size,)
+    assert isinstance(out_size, (tuple, list)), f"Invalid type of out_size: {type(out_size)}"
     self._out_size = tuple(out_size)
 
 
@@ -152,7 +154,8 @@ class Sequential(Module, UpdateReturn, Container, ExplicitInOutSize):
     self.children = visible_module_dict(self.format_elements(object, first, *tuple_modules, **dict_modules))
 
     # the input and output shape
-    self.in_size = tuple(first.in_size)
+    if first.in_size is not None:
+      self.in_size = first.in_size
     self.out_size = tuple(in_size)
 
   def _format_module(self, module, in_size):

brainstate/nn/_dynamics.py

@@ -103,6 +103,9 @@ class IF(Neuron):
   def init_state(self, batch_size: int = None, **kwargs):
     self.V = ShortTermState(init.param(jnp.zeros, self.varshape, batch_size))
 
+  def reset_state(self, batch_size: int = None, **kwargs):
+    self.V.value = init.param(jnp.zeros, self.varshape, batch_size)
+
   def get_spike(self, V=None):
     V = self.V.value if V is None else V
     v_scaled = (V - self.V_th) / self.V_th
@@ -160,6 +163,9 @@ class LIF(Neuron):
   def init_state(self, batch_size: int = None, **kwargs):
     self.V = ShortTermState(init.param(init.Constant(self.V_reset), self.varshape, batch_size))
 
+  def reset_state(self, batch_size: int = None, **kwargs):
+    self.V.value = init.param(init.Constant(self.V_reset), self.varshape, batch_size)
+
   def get_spike(self, V=None):
     V = self.V.value if V is None else V
     v_scaled = (V - self.V_th) / self.V_th
@@ -214,6 +220,10 @@ class ALIF(Neuron):
     self.V = ShortTermState(init.param(init.Constant(0.), self.varshape, batch_size))
     self.a = ShortTermState(init.param(init.Constant(0.), self.varshape, batch_size))
 
+  def reset_state(self, batch_size: int = None, **kwargs):
+    self.V.value = init.param(init.Constant(0.), self.varshape, batch_size)
+    self.a.value = init.param(init.Constant(0.), self.varshape, batch_size)
+
   def get_spike(self, V=None, a=None):
     V = self.V.value if V is None else V
     a = self.a.value if a is None else a
@@ -275,6 +285,9 @@ class Expon(Synapse):
   def init_state(self, batch_size: int = None, **kwargs):
     self.g = ShortTermState(init.param(init.Constant(0.), self.varshape, batch_size))
 
+  def reset_state(self, batch_size: int = None, **kwargs):
+    self.g.value = init.param(init.Constant(0.), self.varshape, batch_size)
+
   def update(self, x=None):
     self.g.value = exp_euler_step(self.dg, self.g.value, environ.get('t'))
     if x is not None:
@@ -325,6 +338,10 @@ class STP(Synapse):
     self.x = ShortTermState(init.param(init.Constant(1.), self.varshape, batch_size))
     self.u = ShortTermState(init.param(init.Constant(self.U), self.varshape, batch_size))
 
+  def reset_state(self, batch_size: int = None, **kwargs):
+    self.x.value = init.param(init.Constant(1.), self.varshape, batch_size)
+    self.u.value = init.param(init.Constant(self.U), self.varshape, batch_size)
+
   def du(self, u, t):
     return self.U - u / self.tau_f
 
@@ -390,6 +407,9 @@ class STD(Synapse):
   def init_state(self, batch_size: int = None, **kwargs):
     self.x = ShortTermState(init.param(init.Constant(1.), self.varshape, batch_size))
 
+  def reset_state(self, batch_size: int = None, **kwargs):
+    self.x.value = init.param(init.Constant(1.), self.varshape, batch_size)
+
   def update(self, pre_spike):
     t = environ.get('t')
     x = exp_euler_step(self.dx, self.x.value, t)

brainstate/nn/_embedding.py

@@ -0,0 +1,66 @@
+# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from typing import Optional, Callable, Union
+
+from ._base import DnnLayer
+from .. import init
+from .._state import ParamState
+from ..mixin import Mode, Training
+from ..typing import ArrayLike
+
+__all__ = [
+  'Embedding',
+]
+
+
+class Embedding(DnnLayer):
+  r"""
+  A simple lookup table that stores embeddings of a fixed size.
+
+  Args:
+    num_embeddings: Size of embedding dictionary. Must be non-negative.
+    embedding_size: Size of each embedding vector. Must be non-negative.
+    embed_init: The initializer for the embedding lookup table, of shape `(num_embeddings, embedding_size)`.
+
+  """
+
+  def __init__(
+      self,
+      num_embeddings: int,
+      embedding_size: int,
+      embed_init: Union[Callable, ArrayLike] = init.LecunUniform(),
+      name: Optional[str] = None,
+      mode: Optional[Mode] = None,
+  ):
+    super().__init__(name=name, mode=mode)
+    if num_embeddings < 0:
+      raise ValueError("num_embeddings must not be negative.")
+    if embedding_size < 0:
+      raise ValueError("embedding_size must not be negative.")
+    self.num_embeddings = num_embeddings
+    self.embedding_size = embedding_size
+    self.out_size = (embedding_size,)
+
+    weight = init.param(embed_init, (self.num_embeddings, self.embedding_size))
+    if self.mode.has(Training):
+      self.weight = ParamState(weight)
+    else:
+      self.weight = weight
+
+  def update(self, indices: ArrayLike):
+    if self.mode.has(Training):
+      return self.weight.value[indices]
+    return self.weight[indices]

brainstate/nn/_rate_rnns.py

@@ -90,6 +90,9 @@ class ValinaRNNCell(RNNCell):
   def init_state(self, batch_size: int = None, **kwargs):
     self.h = ShortTermState(init.param(self._state_initializer, self.num_out, batch_size))
 
+  def reset_state(self, batch_size: int = None, **kwargs):
+    self.h.value = init.param(self._state_initializer, self.num_out, batch_size)
+
   def update(self, x):
     xh = jnp.concatenate([x, self.h.value], axis=-1)
     h = self.W(xh)
@@ -147,6 +150,9 @@ class GRUCell(RNNCell):
   def init_state(self, batch_size: int = None, **kwargs):
     self.h = ShortTermState(init.param(self._state_initializer, [self.num_out], batch_size))
 
+  def reset_state(self, batch_size: int = None, **kwargs):
+    self.h.value = init.param(self._state_initializer, [self.num_out], batch_size)
+
   def update(self, x):
     old_h = self.h.value
     xh = jnp.concatenate([x, old_h], axis=-1)
@@ -224,6 +230,9 @@ class MGUCell(RNNCell):
   def init_state(self, batch_size: int = None, **kwargs):
     self.h = ShortTermState(init.param(self._state_initializer, [self.num_out], batch_size))
 
+  def reset_state(self, batch_size: int = None, **kwargs):
+    self.h.value = init.param(self._state_initializer, [self.num_out], batch_size)
+
   def update(self, x):
     old_h = self.h.value
     xh = jnp.concatenate([x, old_h], axis=-1)
@@ -327,6 +336,10 @@ class LSTMCell(RNNCell):
     self.c = ShortTermState(init.param(self._state_initializer, [self.num_out], batch_size))
     self.h = ShortTermState(init.param(self._state_initializer, [self.num_out], batch_size))
 
+  def reset_state(self, batch_size: int = None, **kwargs):
+    self.c.value = init.param(self._state_initializer, [self.num_out], batch_size)
+    self.h.value = init.param(self._state_initializer, [self.num_out], batch_size)
+
   def update(self, x):
     h, c = self.h.value, self.c.value
     xh = jnp.concat([x, h], axis=-1)
@@ -379,6 +392,10 @@ class URLSTMCell(RNNCell):
     self.c = ShortTermState(init.param(self._state_initializer, [self.num_out], batch_size))
     self.h = ShortTermState(init.param(self._state_initializer, [self.num_out], batch_size))
 
+  def reset_state(self, batch_size: int = None, **kwargs):
+    self.c.value = init.param(self._state_initializer, [self.num_out], batch_size)
+    self.h.value = init.param(self._state_initializer, [self.num_out], batch_size)
+
   def update(self, x: ArrayLike) -> ArrayLike:
     h, c = self.h.value, self.c.value
     xh = jnp.concat([x, h], axis=-1)

brainstate/nn/_readout.py

@@ -66,6 +66,9 @@ class LeakyRateReadout(DnnLayer):
   def init_state(self, batch_size=None, **kwargs):
     self.r = ShortTermState(init.param(init.Constant(0.), self.out_size, batch_size))
 
+  def reset_state(self, batch_size=None, **kwargs):
+    self.r.value = init.param(init.Constant(0.), self.out_size, batch_size)
+
   def update(self, x):
     r = self.decay * self.r.value + x @ self.weight.value
     self.r.value = r
@@ -109,6 +112,9 @@ class LeakySpikeReadout(Neuron):
   def init_state(self, batch_size, **kwargs):
     self.V = ShortTermState(init.param(init.Constant(0.), self.varshape, batch_size))
 
+  def reset_state(self, batch_size, **kwargs):
+    self.V.value = init.param(init.Constant(0.), self.varshape, batch_size)
+
   @property
   def spike(self):
     return self.get_spike(self.V.value)

brainstate/optim/_lr_scheduler_test.py

@@ -34,3 +34,16 @@ class TestMultiStepLR(unittest.TestCase):
         self.assertTrue(jnp.allclose(r, 0.001))
       else:
         self.assertTrue(jnp.allclose(r, 0.0001))
+
+  def test2(self):
+    lr = bst.transform.jit(bst.optim.MultiStepLR(0.1, [10, 20, 30], gamma=0.1))
+    for i in range(40):
+      r = lr(i)
+      if i < 10:
+        self.assertEqual(r, 0.1)
+      elif i < 20:
+        self.assertTrue(jnp.allclose(r, 0.01))
+      elif i < 30:
+        self.assertTrue(jnp.allclose(r, 0.001))
+      else:
+        self.assertTrue(jnp.allclose(r, 0.0001))

brainstate/transform/_jit.py

@@ -23,8 +23,8 @@ import jax
 from jax._src import sharding_impls
 from jax.lib import xla_client as xc
 
-from ._make_jaxpr import StatefulFunction, _ensure_index_tuple, _assign_state_values
 from brainstate._utils import set_module_as
+from ._make_jaxpr import StatefulFunction, _ensure_index_tuple, _assign_state_values
 
 __all__ = ['jit']
 
@@ -33,10 +33,13 @@ class JittedFunction(Callable):
   """
   A wrapped version of ``fun``, set up for just-in-time compilation.
   """
-  origin_fun: Callable   # the original function
+  origin_fun: Callable  # the original function
   stateful_fun: StatefulFunction  # the stateful function for extracting states
   jitted_fun: jax.stages.Wrapped  # the jitted function
-  clear_cache: Callable # clear the cache of the jitted function
+  clear_cache: Callable  # clear the cache of the jitted function
+
+  def __call__(self, *args, **kwargs):
+    pass
 
 
 def _get_jitted_fun(
@@ -85,12 +88,16 @@ def _get_jitted_fun(
     jit_fun.clear_cache()
 
   jitted_fun: JittedFunction
+
   # the original function
   jitted_fun.origin_fun = fun.fun
+
   # the stateful function for extracting states
   jitted_fun.stateful_fun = fun
+
   # the jitted function
   jitted_fun.jitted_fun = jit_fun
+
   # clear cache
   jitted_fun.clear_cache = clear_cache
 
@@ -99,18 +106,18 @@ def _get_jitted_fun(
 
 @set_module_as('brainstate.transform')
 def jit(
-  fun: Callable = None,
-  in_shardings=sharding_impls.UNSPECIFIED,
-  out_shardings=sharding_impls.UNSPECIFIED,
-  static_argnums: int | Sequence[int] | None = None,
-  donate_argnums: int | Sequence[int] | None = None,
-  donate_argnames: str | Iterable[str] | None = None,
-  keep_unused: bool = False,
-  device: xc.Device | None = None,
-  backend: str | None = None,
-  inline: bool = False,
-  abstracted_axes: Any | None = None,
-  **kwargs
+    fun: Callable = None,
+    in_shardings=sharding_impls.UNSPECIFIED,
+    out_shardings=sharding_impls.UNSPECIFIED,
+    static_argnums: int | Sequence[int] | None = None,
+    donate_argnums: int | Sequence[int] | None = None,
+    donate_argnames: str | Iterable[str] | None = None,
+    keep_unused: bool = False,
+    device: xc.Device | None = None,
+    backend: str | None = None,
+    inline: bool = False,
+    abstracted_axes: Any | None = None,
+    **kwargs
 ) -> Union[JittedFunction, Callable[[Callable], JittedFunction]]:
   """
   Sets up ``fun`` for just-in-time compilation with XLA.
@@ -228,12 +235,31 @@ def jit(
 
   if fun is None:
     def wrapper(fun_again: Callable) -> JittedFunction:
-      return _get_jitted_fun(fun_again, in_shardings, out_shardings, static_argnums,
-                             donate_argnums, donate_argnames, keep_unused,
-                             device, backend, inline, abstracted_axes, **kwargs)
+      return _get_jitted_fun(fun_again,
+                             in_shardings,
+                             out_shardings,
+                             static_argnums,
+                             donate_argnums,
+                             donate_argnames,
+                             keep_unused,
+                             device,
+                             backend,
+                             inline,
+                             abstracted_axes,
+                             **kwargs)
+
     return wrapper
 
   else:
-    return _get_jitted_fun(fun, in_shardings, out_shardings, static_argnums,
-                           donate_argnums, donate_argnames, keep_unused,
-                           device, backend, inline, abstracted_axes, **kwargs)
+    return _get_jitted_fun(fun,
+                           in_shardings,
+                           out_shardings,
+                           static_argnums,
+                           donate_argnums,
+                           donate_argnames,
+                           keep_unused,
+                           device,
+                           backend,
+                           inline,
+                           abstracted_axes,
+                           **kwargs)