brainstate 0.0.2.post20241010__py2.py3-none-any.whl → 0.1.0.post20241122__py2.py3-none-any.whl

brainstate/mixin_test.py

@@ -19,57 +19,55 @@ import brainstate as bc
 
 
 class TestMixin(unittest.TestCase):
-  def test_mixin(self):
-    self.assertTrue(bc.mixin.Mixin)
-    self.assertTrue(bc.mixin.DelayedInit)
-    self.assertTrue(bc.mixin.DelayedInitializer)
-    self.assertTrue(bc.mixin.JointTypes)
-    self.assertTrue(bc.mixin.OneOfTypes)
-    self.assertTrue(bc.mixin.Mode)
-    self.assertTrue(bc.mixin.Batching)
-    self.assertTrue(bc.mixin.Training)
-
-
+    def test_mixin(self):
+        self.assertTrue(bc.mixin.Mixin)
+        self.assertTrue(bc.mixin.ParamDesc)
+        self.assertTrue(bc.mixin.ParamDescriber)
+        self.assertTrue(bc.mixin.JointTypes)
+        self.assertTrue(bc.mixin.OneOfTypes)
+        self.assertTrue(bc.mixin.Mode)
+        self.assertTrue(bc.mixin.Batching)
+        self.assertTrue(bc.mixin.Training)
 
 
 class TestMode(unittest.TestCase):
-  def test_JointMode(self):
-    a = bc.mixin.JointMode(bc.mixin.Batching(), bc.mixin.Training())
-    self.assertTrue(a.is_a(bc.mixin.JointTypes[bc.mixin.Batching, bc.mixin.Training]))
-    self.assertTrue(a.has(bc.mixin.Batching))
-    self.assertTrue(a.has(bc.mixin.Training))
-    b = bc.mixin.JointMode(bc.mixin.Batching())
-    self.assertTrue(b.is_a(bc.mixin.JointTypes[bc.mixin.Batching]))
-    self.assertTrue(b.is_a(bc.mixin.Batching))
-    self.assertTrue(b.has(bc.mixin.Batching))
+    def test_JointMode(self):
+        a = bc.mixin.JointMode(bc.mixin.Batching(), bc.mixin.Training())
+        self.assertTrue(a.is_a(bc.mixin.JointTypes[bc.mixin.Batching, bc.mixin.Training]))
+        self.assertTrue(a.has(bc.mixin.Batching))
+        self.assertTrue(a.has(bc.mixin.Training))
+        b = bc.mixin.JointMode(bc.mixin.Batching())
+        self.assertTrue(b.is_a(bc.mixin.JointTypes[bc.mixin.Batching]))
+        self.assertTrue(b.is_a(bc.mixin.Batching))
+        self.assertTrue(b.has(bc.mixin.Batching))
 
-  def test_Training(self):
-    a = bc.mixin.Training()
-    self.assertTrue(a.is_a(bc.mixin.Training))
-    self.assertTrue(a.is_a(bc.mixin.JointTypes[bc.mixin.Training]))
-    self.assertTrue(a.has(bc.mixin.Training))
-    self.assertTrue(a.has(bc.mixin.JointTypes[bc.mixin.Training]))
-    self.assertFalse(a.is_a(bc.mixin.Batching))
-    self.assertFalse(a.has(bc.mixin.Batching))
+    def test_Training(self):
+        a = bc.mixin.Training()
+        self.assertTrue(a.is_a(bc.mixin.Training))
+        self.assertTrue(a.is_a(bc.mixin.JointTypes[bc.mixin.Training]))
+        self.assertTrue(a.has(bc.mixin.Training))
+        self.assertTrue(a.has(bc.mixin.JointTypes[bc.mixin.Training]))
+        self.assertFalse(a.is_a(bc.mixin.Batching))
+        self.assertFalse(a.has(bc.mixin.Batching))
 
-  def test_Batching(self):
-    a = bc.mixin.Batching()
-    self.assertTrue(a.is_a(bc.mixin.Batching))
-    self.assertTrue(a.is_a(bc.mixin.JointTypes[bc.mixin.Batching]))
-    self.assertTrue(a.has(bc.mixin.Batching))
-    self.assertTrue(a.has(bc.mixin.JointTypes[bc.mixin.Batching]))
+    def test_Batching(self):
+        a = bc.mixin.Batching()
+        self.assertTrue(a.is_a(bc.mixin.Batching))
+        self.assertTrue(a.is_a(bc.mixin.JointTypes[bc.mixin.Batching]))
+        self.assertTrue(a.has(bc.mixin.Batching))
+        self.assertTrue(a.has(bc.mixin.JointTypes[bc.mixin.Batching]))
 
-    self.assertFalse(a.is_a(bc.mixin.Training))
-    self.assertFalse(a.has(bc.mixin.Training))
+        self.assertFalse(a.is_a(bc.mixin.Training))
+        self.assertFalse(a.has(bc.mixin.Training))
 
-  def test_Mode(self):
-    a = bc.mixin.Mode()
-    self.assertTrue(a.is_a(bc.mixin.Mode))
-    self.assertTrue(a.is_a(bc.mixin.JointTypes[bc.mixin.Mode]))
-    self.assertTrue(a.has(bc.mixin.Mode))
-    self.assertTrue(a.has(bc.mixin.JointTypes[bc.mixin.Mode]))
+    def test_Mode(self):
+        a = bc.mixin.Mode()
+        self.assertTrue(a.is_a(bc.mixin.Mode))
+        self.assertTrue(a.is_a(bc.mixin.JointTypes[bc.mixin.Mode]))
+        self.assertTrue(a.has(bc.mixin.Mode))
+        self.assertTrue(a.has(bc.mixin.JointTypes[bc.mixin.Mode]))
 
-    self.assertFalse(a.is_a(bc.mixin.Training))
-    self.assertFalse(a.has(bc.mixin.Training))
-    self.assertFalse(a.is_a(bc.mixin.Batching))
-    self.assertFalse(a.has(bc.mixin.Batching))
+        self.assertFalse(a.is_a(bc.mixin.Training))
+        self.assertFalse(a.has(bc.mixin.Training))
+        self.assertFalse(a.is_a(bc.mixin.Batching))
+        self.assertFalse(a.has(bc.mixin.Batching))

brainstate/nn/__init__.py

@@ -13,65 +13,40 @@
 # limitations under the License.
 # ==============================================================================
 
+
 from . import metrics
-from ._base import *
-from ._base import __all__ as base_all
-from ._connections import *
-from ._connections import __all__ as connections_all
+from ._collective_ops import *
+from ._collective_ops import __all__ as collective_ops_all
+from ._dyn_impl import *
+from ._dyn_impl import __all__ as dyn_impl_all
 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 *
-from ._normalizations import __all__ as normalizations_all
-from ._others import *
-from ._others import __all__ as others_all
-from ._poolings import *
-from ._poolings import __all__ as poolings_all
-from ._projection import *
-from ._projection import __all__ as _projection_all
-from ._rate_rnns import *
-from ._rate_rnns import __all__ as rate_rnns
-from ._readout import *
-from ._readout import __all__ as readout_all
-from ._synouts import *
-from ._synouts import __all__ as synouts_all
-from .event import *
-from .event import __all__ as event_all
+from ._exp_euler import *
+from ._exp_euler import __all__ as exp_euler_all
+from ._interaction import *
+from ._interaction import __all__ as interaction_all
+from ._module import *
+from ._module import __all__ as module_all
 
 __all__ = (
-    base_all +
-    connections_all +
-    dynamics_all +
-    elementwise_all +
-    embed_all +
-    normalizations_all +
-    others_all +
-    poolings_all +
-    rate_rnns +
-    readout_all +
-    synouts_all +
-    _projection_all +
-    _misc_all +
-    event_all
+    ['metrics']
+    + collective_ops_all
+    + dyn_impl_all
+    + dynamics_all
+    + elementwise_all
+    + module_all
+    + exp_euler_all
+    + interaction_all
 )
 
 del (
-  base_all,
-  connections_all,
-  dynamics_all,
-  elementwise_all,
-  embed_all,
-  normalizations_all,
-  others_all,
-  poolings_all,
-  readout_all,
-  synouts_all,
-  _projection_all,
-  _misc_all,
-  event_all
+    collective_ops_all,
+    dyn_impl_all,
+    dynamics_all,
+    elementwise_all,
+    module_all,
+    exp_euler_all,
+    interaction_all,
 )

brainstate/nn/_collective_ops.py

@@ -0,0 +1,199 @@
+# 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 collections import namedtuple
+from typing import Dict, Callable, TypeVar
+
+import jax
+
+from brainstate._utils import set_module_as
+from brainstate.graph import nodes
+from ._module import Module
+
+# the maximum order
+MAX_ORDER = 10
+
+# State Load Results
+StateLoadResult = namedtuple('StateLoadResult', ['missing_keys', 'unexpected_keys'])
+
+T = TypeVar('T', bound=Module)
+
+__all__ = [
+    'MAX_ORDER', 'call_order', 'init_all_states', 'reset_all_states',
+    'load_all_states', 'save_all_states', 'assign_state_values',
+]
+
+
+@set_module_as('brainstate.nn')
+def call_order(level: int = 0, check_order_boundary: bool = True):
+    """The decorator for indicating the resetting level.
+
+    The function takes an optional integer argument level with a default value of 0.
+
+    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)
+
+    Parameters
+    ----------
+    level: int
+      The call order level.
+    check_order_boundary: bool
+      Whether check the boundary of function call order. If True,
+      the order that not in [0,  10) will raise a ValueError.
+
+    Returns
+    -------
+    The function to warp.
+    """
+    if check_order_boundary and (level < 0 or level >= MAX_ORDER):
+        raise ValueError(f'"call_order" must be an integer in [0, {MAX_ORDER}). but we got {level}.')
+
+    def wrap(fun: Callable):
+        fun.call_order = level
+        return fun
+
+    return wrap
+
+
+@set_module_as('brainstate.nn')
+def init_all_states(target: T, *args, exclude=None, **kwargs) -> T:
+    """
+    Collectively initialize states of all children nodes in the given target.
+
+    Args:
+      target: The target Module.
+
+    Returns:
+      The target Module.
+    """
+    nodes_with_order = []
+
+    nodes_ = nodes(target).filter(Module)
+    if exclude is not None:
+        nodes_ = nodes_ - nodes_.filter(exclude)
+
+    # reset node whose `init_state` has no `call_order`
+    for node in list(nodes_.values()):
+        if hasattr(node.init_state, 'call_order'):
+            nodes_with_order.append(node)
+        else:
+            node.init_state(*args, **kwargs)
+
+    # reset the node's states
+    for node in sorted(nodes_with_order, key=lambda x: x.init_state.call_order):
+        node.init_state(*args, **kwargs)
+
+    return target
+
+
+@set_module_as('brainstate.nn')
+def reset_all_states(target: Module, *args, **kwargs) -> Module:
+    """
+    Collectively 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 path, node in nodes(target).filter(Module).items():
+        if hasattr(node.reset_state, 'call_order'):
+            nodes_with_order.append(node)
+        else:
+            node.reset_state(*args, **kwargs)
+
+    # 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.nn')
+def load_all_states(target: Module, state_dict: Dict, **kwargs):
+    """
+    Copy parameters and buffers from :attr:`state_dict` into
+    this module and its descendants.
+
+    Args:
+      target: Module. The dynamical system to load its states.
+      state_dict: dict. A dict containing parameters and persistent buffers.
+
+    Returns
+    -------
+      ``NamedTuple``  with ``missing_keys`` and ``unexpected_keys`` fields:
+
+      * **missing_keys** is a list of str containing the missing keys
+      * **unexpected_keys** is a list of str containing the unexpected keys
+    """
+    missing_keys = []
+    unexpected_keys = []
+    for path, node in nodes(target).items():
+        r = node.load_state(state_dict[path], **kwargs)
+        if r is not None:
+            missing, unexpected = r
+            missing_keys.extend([f'{path}.{key}' for key in missing])
+            unexpected_keys.extend([f'{path}.{key}' for key in unexpected])
+    return StateLoadResult(missing_keys, unexpected_keys)
+
+
+@set_module_as('brainstate.nn')
+def save_all_states(target: Module, **kwargs) -> Dict:
+    """
+    Save all states in the ``target`` as a dictionary for later disk serialization.
+
+    Args:
+      target: Module. The node to save its states.
+
+    Returns:
+      Dict. The state dict for serialization.
+    """
+    return {key: node.save_state(**kwargs) for key, node in target.nodes().items()}
+
+
+@set_module_as('brainstate.nn')
+def assign_state_values(target: Module, *state_by_abs_path: Dict):
+    """
+    Assign state values according to the given state dictionary.
+
+    Parameters
+    ----------
+    target: Module
+      The target module.
+    state_by_abs_path: dict
+      The state dictionary which is accessed by the "absolute" accessing method.
+
+    """
+    all_states = dict()
+    for state in state_by_abs_path:
+        all_states.update(state)
+    variables = target.states()
+    keys1 = set(all_states.keys())
+    keys2 = set(variables.keys())
+    for key in keys2.intersection(keys1):
+        variables[key].value = jax.numpy.asarray(all_states[key])
+    unexpected_keys = list(keys1 - keys2)
+    missing_keys = list(keys2 - keys1)
+    return unexpected_keys, missing_keys

brainstate/nn/_dyn_impl/__init__.py

@@ -0,0 +1,46 @@
+# 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 ._dynamics_neuron import *
+from ._dynamics_neuron import __all__ as dyn_neuron_all
+from ._dynamics_synapse import *
+from ._dynamics_synapse import __all__ as dyn_synapse_all
+from ._inputs import *
+from ._inputs import __all__ as inputs_all
+from ._projection_alignpost import *
+from ._projection_alignpost import __all__ as alignpost_all
+from ._rate_rnns import *
+from ._rate_rnns import __all__ as rate_rnns
+from ._readout import *
+from ._readout import __all__ as readout_all
+
+__all__ = (
+    dyn_neuron_all
+    + dyn_synapse_all
+    + inputs_all
+    + alignpost_all
+    + rate_rnns
+    + readout_all
+)
+
+del (
+    dyn_neuron_all,
+    dyn_synapse_all,
+    inputs_all,
+    readout_all,
+    alignpost_all,
+    rate_rnns,
+)