braintrace 0.1.1__py2.py3-none-any.whl

braintrace/__init__.py

@@ -0,0 +1,79 @@
+# Copyright 2024 BrainX 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.
+# ==============================================================================
+
+# -*- coding: utf-8 -*-
+
+
+__version__ = "0.1.1"
+__version_info__ = (0, 1, 1)
+
+from braintrace._etrace_algorithms import *
+from braintrace._etrace_algorithms import __all__ as _alg_all
+from braintrace._etrace_compiler_graph import *
+from braintrace._etrace_compiler_graph import __all__ as _compiler_all
+from braintrace._etrace_compiler_hid_param_op import *
+from braintrace._etrace_compiler_hid_param_op import __all__ as _hid_param_all
+from braintrace._etrace_compiler_hidden_group import *
+from braintrace._etrace_compiler_hidden_group import __all__ as _hid_group_all
+from braintrace._etrace_compiler_hidden_pertubation import *
+from braintrace._etrace_compiler_hidden_pertubation import __all__ as _hid_pertub_all
+from braintrace._etrace_compiler_module_info import *
+from braintrace._etrace_compiler_module_info import __all__ as _mod_info_all
+from braintrace._etrace_concepts import *
+from braintrace._etrace_concepts import __all__ as _con_all
+from braintrace._etrace_graph_executor import *
+from braintrace._etrace_graph_executor import __all__ as _exec_all
+from braintrace._etrace_input_data import *
+from braintrace._etrace_input_data import __all__ as _data_all
+from braintrace._etrace_operators import *
+from braintrace._etrace_operators import __all__ as _op_all
+from braintrace._etrace_vjp import *
+from braintrace._etrace_vjp import __all__ as _vjp_all
+from braintrace._grad_exponential import *
+from braintrace._grad_exponential import __all__ as _grad_exp_all
+from braintrace._misc import *
+from braintrace._misc import __all__ as _misc_all
+from . import nn
+
+__all__ = ['nn'] + _alg_all + _compiler_all + _hid_param_all + _hid_group_all + _hid_pertub_all
+__all__ += _mod_info_all + _con_all + _exec_all + _data_all + _op_all + _vjp_all
+__all__ += _grad_exp_all + _misc_all
+
+del _alg_all, _compiler_all, _hid_param_all, _hid_group_all, _hid_pertub_all
+del _mod_info_all, _con_all, _exec_all, _data_all, _op_all, _vjp_all
+del _grad_exp_all,
+del _misc_all
+
+
+def __getattr__(name):
+    mapping = {
+        'ETraceState': 'HiddenState',
+        'ETraceGroupState': 'HiddenGroupState',
+        'ETraceTreeState': 'HiddenTreeState',
+    }
+
+    if name in mapping:
+        import warnings
+        import brainstate
+
+        warnings.warn(
+            f"braintrace.{name} is deprecated and will be removed in a future release. "
+            f"Please use brainstate.{mapping[name]} instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
+        return getattr(brainstate, mapping[name])
+    raise AttributeError(name)

braintrace/_compatible_imports.py

@@ -0,0 +1,62 @@
+# Copyright 2024 BrainX 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.
+# ==============================================================================
+
+import jax
+
+__all__ = [
+    'Primitive',
+    'Var',
+    'JaxprEqn',
+    'Jaxpr',
+    'ClosedJaxpr',
+    'Literal',
+    'new_var',
+    'is_jit_primitive',
+    'is_scan_primitive',
+    'is_while_primitive',
+    'is_cond_primitive',
+]
+
+if jax.__version_info__ < (0, 4, 38):
+    from jax.core import Primitive, Var, JaxprEqn, Jaxpr, ClosedJaxpr, Literal
+
+else:
+    from jax.extend.core import Primitive, Var, JaxprEqn, Jaxpr, ClosedJaxpr, Literal
+
+
+def new_var(suffix, aval):
+    if jax.__version_info__ < (0, 6, 2):
+        return Var(suffix, aval)
+    else:
+        return Var(aval)
+
+
+def is_jit_primitive(eqn: JaxprEqn) -> bool:
+    if jax.__version_info__ < (0, 7, 0):
+        return eqn.primitive.name == 'pjit'
+    else:
+        return eqn.primitive.name == 'jit'
+
+
+def is_scan_primitive(eqn: JaxprEqn) -> bool:
+    return eqn.primitive.name == 'scan'
+
+
+def is_while_primitive(eqn: JaxprEqn) -> bool:
+    return eqn.primitive.name == 'while'
+
+
+def is_cond_primitive(eqn: JaxprEqn) -> bool:
+    return eqn.primitive.name == 'cond'

braintrace/_compatible_imports_test.py

@@ -0,0 +1,94 @@
+# Copyright 2024 BrainX 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.
+# ==============================================================================
+
+import jax.numpy as jnp
+from jax import jit, make_jaxpr, lax
+
+from braintrace._compatible_imports import (
+    is_jit_primitive, is_scan_primitive, is_while_primitive,
+    is_cond_primitive
+)
+
+
+class TestPrimitive:
+    def test_jit(self):
+        @jit
+        def jit_function(x, y):
+            return x ** 2 + jnp.sin(y)
+
+        # Note: make_jaxpr on a jitted function shows the same jaxpr
+        jaxpr_jit = make_jaxpr(jit_function)(2.0, 1.0)
+        assert is_jit_primitive(jaxpr_jit.eqns[0])
+
+    def test_scan(self):
+        print("3. make_jaxpr with lax.scan:")
+
+        def scan_step(carry, x):
+            return carry + x, carry * x
+
+        def scan_function(init, xs):
+            return lax.scan(scan_step, init, xs)
+
+        # Create sample data
+        init_val = 1.0
+        xs = jnp.array([1.0, 2.0, 3.0, 4.0])
+
+        jaxpr_scan = make_jaxpr(scan_function)(init_val, xs)
+        assert is_scan_primitive(jaxpr_scan.eqns[0])
+
+    def test_while(self):
+        def while_cond(carry):
+            i, x = carry
+            return i < 5
+
+        def while_body(carry):
+            i, x = carry
+            return i + 1, x * 2
+
+        def while_function(init_carry):
+            return lax.while_loop(while_cond, while_body, init_carry)
+
+        init_carry = (0, 1.0)
+        jaxpr_while = make_jaxpr(while_function)(init_carry)
+        assert is_while_primitive(jaxpr_while.eqns[0])
+
+    def test_cond(self):
+        def true_branch(x):
+            return x * 2
+
+        def false_branch(x):
+            return x + 1
+
+        def cond_function(pred, x):
+            return lax.cond(pred, true_branch, false_branch, x)
+
+        jaxpr_cond = make_jaxpr(cond_function)(True, 5.0)
+        assert is_cond_primitive(jaxpr_cond.eqns[-1])
+
+    def test_fori_loop(self):
+        def branch_0(x):
+            return x * 2
+
+        def branch_1(x):
+            return x + 10
+
+        def branch_2(x):
+            return x ** 2
+
+        def switch_function(index, x):
+            return lax.switch(index, [branch_0, branch_1, branch_2], x)
+
+        jaxpr_switch = make_jaxpr(switch_function)(1, 3.0)
+        assert is_cond_primitive(jaxpr_switch.eqns[-1])

braintrace/_etrace_algorithms.py

@@ -0,0 +1,333 @@
+# Copyright 2024 BrainX 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.
+#
+# Author: Chaoming Wang <chao.brain@qq.com>
+# Date: 2024-04-03
+# Copyright: 2024, Chaoming Wang
+# ==============================================================================
+
+# -*- coding: utf-8 -*-
+
+from typing import Dict, Any, Optional
+
+import brainstate
+
+from ._etrace_compiler_graph import ETraceGraph
+from ._etrace_graph_executor import ETraceGraphExecutor
+from ._typing import Path
+
+__all__ = [
+    'ETraceAlgorithm',
+    'EligibilityTrace',
+]
+
+
+class EligibilityTrace(brainstate.ShortTermState):
+    """
+    The state for storing the eligibility trace during the computation of
+    online learning algorithms.
+
+    Examples
+    --------
+    When you are using :class:`braintrace.IODimVjpAlgorithm`, you can get
+    the eligibility trace of the weight by calling:
+
+    .. code-block:: python
+
+        >>> etrace = etrace_algorithm.etrace_of(weight)
+
+    """
+    __module__ = 'braintrace'
+
+
+class ETraceAlgorithm(brainstate.nn.Module):
+    r"""
+    The base class for the eligibility trace algorithm.
+
+    Parameters
+    ----------
+    model : brainstate.nn.Module
+        The model function, which receives the input arguments and returns the model output.
+    name : str, optional
+        The name of the etrace algorithm.
+
+    Attributes
+    ----------
+    graph : ETraceGraphExecutor
+        The etrace graph.
+    param_states : Dict[Hashable, brainstate.ParamState]
+        The weight states.
+    hidden_states : Dict[Hashable, brainstate.HiddenState]
+        The hidden states.
+    other_states : Dict[Hashable, brainstate.State]
+        The other states.
+    is_compiled : bool
+        Whether the etrace algorithm has been compiled.
+    running_index : brainstate.ParamState[int]
+        The running index.
+    """
+    __module__ = 'braintrace'
+
+    def __init__(
+        self,
+        model: brainstate.nn.Module,
+        graph_executor: ETraceGraphExecutor,
+        name: Optional[str] = None,
+    ):
+        super().__init__(name=name)
+
+        # the model
+        if not isinstance(model, brainstate.nn.Module):
+            raise ValueError(
+                f'The model should be a brainstate.nn.Module, this can help us to '
+                f'better obtain the program structure. But we got {type(model)}.'
+            )
+        self.model4compile = model
+
+        # the graph
+        if not isinstance(graph_executor, ETraceGraphExecutor):
+            raise ValueError(
+                f'The graph should be a ETraceGraphExecutor, this can help us to '
+                f'better obtain the program structure. But we got {type(graph_executor)}.'
+            )
+        self.graph_executor = graph_executor
+
+        # The flag to indicate whether the etrace algorithm has been compiled
+        self.is_compiled = False
+
+        # the running index
+        self.running_index = brainstate.LongTermState(0)
+
+        # other states
+        self._param_states = None
+        self._hidden_states = None
+        self._other_states = None
+
+    @property
+    def graph(self) -> ETraceGraph:
+        """
+        Get the etrace graph.
+
+        Returns
+        -------
+        ETraceGraph
+            The etrace graph.
+        """
+        return self.graph_executor.graph
+
+    @property
+    def executor(self) -> ETraceGraphExecutor:
+        """
+        Get the etrace graph executor.
+
+        Returns
+        -------
+        ETraceGraphExecutor
+            The etrace graph executor.
+        """
+        return self.graph_executor
+
+    @property
+    def param_states(self) -> brainstate.util.FlattedDict[Path, brainstate.ParamState]:
+        """
+        Get the parameter weight states.
+
+        Returns
+        -------
+        brainstate.util.FlattedDict[Path, brainstate.ParamState]
+            The parameter weight states.
+        """
+        if self._param_states is None:
+            self._split_state()
+        return self._param_states
+
+    @property
+    def hidden_states(self) -> brainstate.util.FlattedDict[Path, brainstate.HiddenState]:
+        """
+        Get the hidden states.
+
+        Returns
+        -------
+        brainstate.util.FlattedDict[Path, brainstate.HiddenState]
+            The hidden states.
+        """
+        if self._hidden_states is None:
+            self._split_state()
+        return self._hidden_states
+
+    @property
+    def other_states(self) -> brainstate.util.FlattedDict[Path, brainstate.State]:
+        """
+        Get the other states.
+
+        Returns
+        -------
+        brainstate.util.FlattedDict[Path, brainstate.State]
+            The other states.
+        """
+        if self._other_states is None:
+            self._split_state()
+        return self._other_states
+
+    def _split_state(self):
+        # --- the state separation --- #
+        #
+        # [NOTE]
+        #
+        # The `ETraceGraphExecutor` and the following states suggests that
+        # `ETraceAlgorithm` depends on the states we created in the
+        # `ETraceGraphExecutor`, including:
+        #
+        #   - the weight states, which is invariant during the training process
+        #   - the hidden states, the recurrent states, which may be changed between different training epochs
+        #   - the other states, which may be changed between different training epochs
+        (
+            self._param_states,
+            self._hidden_states,
+            self._other_states
+        ) = self.graph.module_info.retrieved_model_states.split(brainstate.ParamState, brainstate.HiddenState, ...)
+
+    def compile_graph(self, *args) -> None:
+        r"""
+        Compile the eligibility trace graph of the relationship between etrace weights, states and operators.
+
+        The compilation process includes:
+
+        - building the etrace graph
+        - separating the states
+        - initializing the etrace states
+
+        Parameters
+        ----------
+        *args
+            The input arguments.
+        """
+
+        if not self.is_compiled:
+            # --- the model etrace graph -- #
+            self.graph_executor.compile_graph(*args)
+
+            # --- the initialization of the states --- #
+            self.init_etrace_state(*args)
+
+            # mark the graph is compiled
+            self.is_compiled = True
+
+    @property
+    def path_to_states(self) -> brainstate.util.FlattedDict[Path, brainstate.State]:
+        """
+        Get the path to the states.
+
+        Returns
+        -------
+        brainstate.util.FlattedDict[Path, brainstate.State]
+            The mapping from path to states.
+        """
+        return self.graph_executor.path_to_states
+
+    @property
+    def state_id_to_path(self) -> Dict[int, Path]:
+        """
+        Get the state ID to the path.
+
+        Returns
+        -------
+        Dict[int, Path]
+            The mapping from state ID to path.
+        """
+        return self.graph_executor.state_id_to_path
+
+    def show_graph(self) -> None:
+        """
+        Show the etrace graph.
+        """
+        return self.graph_executor.show_graph()
+
+    def __call__(self, *args) -> Any:
+        """
+        Update the model and the eligibility trace states.
+
+        Parameters
+        ----------
+        *args
+            The input arguments.
+
+        Returns
+        -------
+        Any
+            The output of the update method.
+        """
+        return self.update(*args)
+
+    def update(self, *args) -> Any:
+        """
+        Update the model and the eligibility trace states.
+
+        Parameters
+        ----------
+        *args
+            The input arguments.
+
+        Returns
+        -------
+        Any
+            The model output.
+
+        Raises
+        ------
+        NotImplementedError
+            This method must be implemented by subclasses.
+        """
+        raise NotImplementedError
+
+    def init_etrace_state(self, *args, **kwargs) -> None:
+        """
+        Initialize the eligibility trace states of the etrace algorithm.
+
+        This method is needed after compiling the etrace graph. See `.compile_graph()` for the details.
+
+        Parameters
+        ----------
+        *args
+            The positional arguments.
+        **kwargs
+            The keyword arguments.
+
+        Raises
+        ------
+        NotImplementedError
+            This method must be implemented by subclasses.
+        """
+        raise NotImplementedError
+
+    def get_etrace_of(self, weight: brainstate.ParamState | Path) -> Any:
+        """
+        Get the eligibility trace of the given weight.
+
+        Parameters
+        ----------
+        weight : brainstate.ParamState | Path
+            The parameter weight or path to the weight.
+
+        Returns
+        -------
+        Any
+            The eligibility trace.
+
+        Raises
+        ------
+        NotImplementedError
+            This method must be implemented by subclasses.
+        """
+        raise NotImplementedError