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

brainstate/graph/_graph_convert.py

@@ -0,0 +1,246 @@
+# The file is adapted from the Flax library (https://github.com/google/flax).
+# The credit should go to the Flax authors.
+#
+# Copyright 2024 The Flax Authors & 2024 BDP Ecosystem.
+#
+# 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 typing import Any, Callable, Iterable, TypeVar, Hashable, Optional, Tuple, List
+
+import jax
+
+from brainstate._state import State
+from brainstate.typing import Missing, PyTree, PathParts
+from brainstate.util import PyTreeNode, field
+from ._graph_context import SplitContext, MergeContext, split_context, merge_context
+from ._graph_operation import (RefMap, iter_leaf, _is_graph_node, GraphDef, GraphStateMapping)
+
+__all__ = [
+    'graph_to_tree', 'tree_to_graph', 'NodeStates'
+]
+
+Node = TypeVar('Node')
+Leaf = TypeVar('Leaf')
+
+KeyEntry = TypeVar('KeyEntry', bound=Hashable)
+KeyPath = tuple[KeyEntry, ...]
+Prefix = Any
+
+
+def check_consistent_aliasing(
+    node: Tuple[Any, ...],
+    prefix: Tuple[Any, ...],
+    /,
+    *,
+    node_prefixes: Optional[RefMap[Any, List[Tuple[PathParts, Any]]]] = None,
+):
+    node_prefixes = RefMap() if node_prefixes is None else node_prefixes
+
+    # collect all paths and prefixes for each node
+    for path, value in iter_leaf(node):
+        if isinstance(value, State):
+            # value.check_valid_trace(lambda: f'Trying to extract graph node '
+            #                                 f'from different trace level, got {value!r}')
+            if value in node_prefixes:
+                paths_prefixes = node_prefixes[value]
+                paths_prefixes.append((path, prefix))
+            else:
+                node_prefixes[value] = [(path, prefix)]
+
+    # check for inconsistent aliasing
+    node_msgs = []
+    for node, paths_prefixes in node_prefixes.items():
+        unique_prefixes = {prefix for _, prefix in paths_prefixes}
+        if len(unique_prefixes) > 1:
+            path_prefix_repr = '\n'.join([f'  {"/".join(map(str, path)) if path else "<root>"}: {prefix}'
+                                          for path, prefix in paths_prefixes])
+            nodes_msg = f'Node: {type(node)}\n{path_prefix_repr}'
+            node_msgs.append(nodes_msg)
+
+    if node_msgs:
+        raise ValueError('Inconsistent aliasing detected. The '
+                         'following nodes have different prefixes:\n'
+                         + '\n'.join(node_msgs))
+
+
+# -----------------------------
+# to_tree/from_tree
+# -----------------------------
+
+def broadcast_prefix(
+    prefix_tree: Any,
+    full_tree: Any,
+    prefix_is_leaf: Optional[Callable[[Any], bool]] = None,
+    tree_is_leaf: Optional[Callable[[Any], bool]] = None,
+) -> List[Any]:
+    """
+    Broadcasts a prefix tree to a full tree.
+
+    Args:
+      prefix_tree: A prefix tree.
+      full_tree: A full tree.
+      prefix_is_leaf: A function that checks if a prefix is a leaf.
+      tree_is_leaf: A function that checks if a tree is a leaf.
+
+    Returns:
+      A list of prefixes.
+    """
+    # If prefix_tree is not a tree prefix of full_tree, this code can raise a
+    # ValueError; use prefix_errors to find disagreements and raise more precise
+    # error messages.
+    result = []
+    num_leaves = lambda t: jax.tree_util.tree_structure(t, is_leaf=tree_is_leaf).num_leaves
+    add_leaves = lambda x, subtree: result.extend([x] * num_leaves(subtree))
+    jax.tree.map(add_leaves, prefix_tree, full_tree, is_leaf=prefix_is_leaf)
+    return result
+
+
+class NodeStates(PyTreeNode):
+    _graphdef: GraphDef[Any] | None
+    states: tuple[GraphStateMapping, ...]
+    metadata: Any = field(pytree_node=False)
+
+    @property
+    def graphdef(self) -> GraphDef[Any]:
+        if self._graphdef is None:
+            raise ValueError('No graphdef available')
+        return self._graphdef
+
+    @property
+    def state(self) -> GraphStateMapping:
+        if len(self.states) != 1:
+            raise ValueError(f'Expected exactly one GraphDefState, got {len(self.states)}')
+        return self.states[0]
+
+    @classmethod
+    def from_split(
+        cls,
+        graphdef: GraphDef[Any],
+        state: GraphStateMapping,
+        /,
+        *states: GraphStateMapping,
+        metadata: Any = None,
+    ):
+        return cls(_graphdef=graphdef, states=(state, *states), metadata=metadata)
+
+    @classmethod
+    def from_states(cls, state: GraphStateMapping, *states: GraphStateMapping):
+        return cls(_graphdef=None, states=(state, *states), metadata=None)
+
+    @classmethod
+    def from_prefixes(cls, prefixes: Iterable[Any], /, *, metadata: Any = None):
+        return cls(_graphdef=None, states=tuple(prefixes), metadata=metadata)
+
+
+def _default_split_fn(ctx: SplitContext, path: KeyPath, prefix: Prefix, leaf: Leaf):
+    return NodeStates.from_split(*ctx.treefy_split(leaf))
+
+
+def graph_to_tree(
+    may_have_graph_nodes,
+    /,
+    *,
+    prefix: Any = Missing,
+    split_fn: Callable[[SplitContext, KeyPath, Prefix, Leaf], Any] = _default_split_fn,
+    map_non_graph_nodes: bool = False,
+    ctxtag: str | None = None,
+    check_aliasing: bool = True,
+) -> PyTree:
+    """
+    Convert a tree of pytree objects to a tree of TreeNode objects.
+    """
+    leaf_prefixes = broadcast_prefix(prefix, may_have_graph_nodes, prefix_is_leaf=lambda x: x is None)
+    leaf_keys, treedef = jax.tree_util.tree_flatten_with_path(may_have_graph_nodes)
+
+    # Check that the number of keys and prefixes match
+    assert len(leaf_keys) == len(leaf_prefixes)
+
+    # Split the tree
+    with split_context(ctxtag) as ctx:
+        leaves_out = []
+        node_prefixes = RefMap[Any, list[tuple[PathParts, Any]]]()
+        for (keypath, leaf), leaf_prefix in zip(leaf_keys, leaf_prefixes):
+            if _is_graph_node(leaf):
+                if check_aliasing:
+                    check_consistent_aliasing(leaf, leaf_prefix, node_prefixes=node_prefixes)
+                leaves_out.append(split_fn(ctx, keypath, leaf_prefix, leaf))
+            else:
+                if map_non_graph_nodes:
+                    leaf = split_fn(ctx, keypath, leaf_prefix, leaf)
+                leaves_out.append(leaf)
+
+    pytree_out = jax.tree.unflatten(treedef, leaves_out)
+    return pytree_out
+
+
+def _is_tree_node(x):
+    """Check if x is a TreeNode."""
+    return isinstance(x, NodeStates)
+
+
+def _merge_tree_node(ctx: MergeContext, path: KeyPath, prefix: Prefix, leaf: Leaf) -> Any:
+    if not isinstance(leaf, NodeStates):
+        raise ValueError(f'Expected TreeNode, got {type(leaf)} at path {path}')
+    return ctx.treefy_merge(leaf.graphdef, *leaf.states)
+
+
+def tree_to_graph(
+    tree: Any,
+    /,
+    *,
+    prefix: Any = Missing,
+    merge_fn: Callable[[MergeContext, KeyPath, Prefix, Leaf], Any] = _merge_tree_node,
+    is_node_leaf: Callable[[Leaf], bool] = _is_tree_node,
+    is_leaf: Callable[[Leaf], bool] = _is_tree_node,
+    map_non_graph_nodes: bool = False,
+    ctxtag: str | None = None,
+) -> Any:
+    """
+    Convert a tree of TreeNode objects to a tree of pytree objects.
+
+    Args:
+      tree: A tree of TreeNode objects.
+      prefix: A tree of prefixes.
+      merge_fn: A function that merges a TreeNode object.
+      is_node_leaf: A function that checks if a leaf is a TreeNode.
+      is_leaf: A function that checks if a leaf is a TreeNode.
+      map_non_graph_nodes: A boolean indicating whether to map non-graph nodes.
+
+    Returns:
+      A tree of pytree objects.
+    """
+    _prefix_is_leaf = lambda x: x is None or is_leaf(x)
+    leaf_prefixes = broadcast_prefix(prefix, tree, prefix_is_leaf=_prefix_is_leaf, tree_is_leaf=is_leaf)
+    leaf_keys, treedef = jax.tree_util.tree_flatten_with_path(tree, is_leaf=is_leaf)
+    assert len(leaf_keys) == len(leaf_prefixes), "Mismatched number of keys and prefixes"
+
+    with merge_context(ctxtag) as ctx:
+        leaves_out = []
+        for (keypath, leaf), leaf_prefix in zip(leaf_keys, leaf_prefixes):
+            if is_node_leaf(leaf):
+                leaf_out = merge_fn(ctx, keypath, leaf_prefix, leaf)
+                leaves_out.append(leaf_out)
+            else:
+                if map_non_graph_nodes:
+                    leaf = merge_fn(ctx, keypath, leaf_prefix, leaf)
+                leaves_out.append(leaf)
+
+    pytree_out = jax.tree.unflatten(treedef, leaves_out)
+    return pytree_out
+
+
+def clear_non_graph_nodes(tree):
+    return jax.tree.map(lambda x: x if _is_graph_node(x) else None, tree)

brainstate/graph/_graph_node.py

@@ -0,0 +1,300 @@
+# The file is adapted from the Flax library (https://github.com/google/flax).
+# The credit should go to the Flax authors.
+#
+# Copyright 2024 The Flax Authors & 2024 BDP Ecosystem.
+#
+# 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 abc import ABCMeta
+from copy import deepcopy
+from typing import Any, Callable, Type, TypeVar, Tuple, TYPE_CHECKING, Mapping, Iterator, Sequence
+
+import brainunit as u
+import jax
+import numpy as np
+
+from brainstate._state import State, TreefyState
+from brainstate.typing import Key
+from brainstate.util._error import TraceContextError
+from brainstate.util._pretty_repr import PrettyRepr, pretty_repr_avoid_duplicate, PrettyType, PrettyAttr
+from brainstate.util._tracers import StateJaxTracer
+from ._graph_operation import register_graph_node_type
+
+__all__ = [
+    'Node', 'Dict', 'List', 'Sequential',
+]
+
+G = TypeVar('G', bound='Node')
+A = TypeVar('A')
+
+
+class GraphNodeMeta(ABCMeta):
+    if not TYPE_CHECKING:
+        def __call__(cls, *args: Any, **kwargs: Any) -> Any:
+            node = cls.__new__(cls, *args, **kwargs)
+            vars(node)['_trace_state'] = StateJaxTracer()
+            node.__init__(*args, **kwargs)
+            return node
+
+
+class Node(PrettyRepr, metaclass=GraphNodeMeta):
+    """
+    Base class for all graph nodes.
+
+    This class provides the following functionalities:
+    - Register the node type with the graph tool.
+    - Prevent mutation of the node from different trace level.
+    - Provide a pretty repr for the node.
+    - Provide a treescope repr for the node.
+    - Deepcopy the node.
+
+    """
+    if TYPE_CHECKING:
+        _trace_state: StateJaxTracer
+
+    def __init_subclass__(cls) -> None:
+        super().__init_subclass__()
+
+        register_graph_node_type(
+            type=cls,
+            flatten=_node_flatten,
+            set_key=_node_set_key,
+            pop_key=_node_pop_key,
+            create_empty=_node_create_empty,
+            clear=_node_clear,
+        )
+
+    # if not TYPE_CHECKING:
+    #   def __setattr__(self, name: str, value: Any) -> None:
+    #     self._setattr(name, value)
+
+    # def _setattr(self, name: str, value: Any) -> None:
+    #   self.check_valid_context(lambda: f"Cannot mutate '{type(self).__name__}' from different trace level")
+    #   object.__setattr__(self, name, value)
+
+    def check_valid_context(self, error_msg: Callable[[], str]) -> None:
+        """
+        Check if the current context is valid for the object to be mutated.
+        """
+        if not self._trace_state.is_valid():
+            raise TraceContextError(error_msg())
+
+    def __deepcopy__(self: G, memo=None) -> G:
+        """
+        Deepcopy the object.
+        """
+        from ._graph_operation import treefy_split, treefy_merge
+
+        graphdef, state = treefy_split(self)
+        graphdef = deepcopy(graphdef)
+        state = deepcopy(state)
+        return treefy_merge(graphdef, state)
+
+    def __pretty_repr__(self):
+        """
+        Pretty repr for the object.
+        """
+        yield from pretty_repr_avoid_duplicate(self, _default_repr_object, _default_repr_attr)
+
+    def __treescope_repr__(self, path, subtree_renderer):
+        """
+        Treescope repr for the object.
+        """
+        children = {}
+        for name, value in vars(self).items():
+            name, value = self.__leaf_fn__(name, value)
+            if name.startswith('_'):
+                continue
+            children[name] = value
+        import treescope  # type: ignore[import-not-found,import-untyped]
+        return treescope.repr_lib.render_object_constructor(
+            object_type=type(self),
+            attributes=children,
+            path=path,
+            subtree_renderer=subtree_renderer,
+            color=treescope.formatting_util.color_from_string(type(self).__qualname__)
+        )
+
+    def __leaf_fn__(self, leaf, value):
+        return leaf, value
+
+
+def _default_repr_object(node: Node):
+    yield PrettyType(type=type(node))
+
+
+def _default_repr_attr(node: Node):
+    for name, value in vars(node).items():
+        name, value = node.__leaf_fn__(name, value)
+        if name.startswith('_'):
+            continue
+        # value = jax.tree.map(_to_shape_dtype, value, is_leaf=lambda x: isinstance(x, u.Quantity))
+        yield PrettyAttr(name, repr(value))
+
+
+class String:
+    def __init__(self, msg):
+        self.msg = msg
+
+    def __repr__(self):
+        return self.msg
+
+
+def _to_shape_dtype(value):
+    if isinstance(value, State):
+        return value.replace(jax.tree.map(_to_shape_dtype, value.value))
+    elif isinstance(value, (np.ndarray, jax.Array)):
+        return String(f'Array(shape={value.shape}, dtype={value.dtype.name})')
+    elif isinstance(value, u.Quantity):
+        return String(f'Quantity(mantissa=Array(shape={value.shape}, dtype={value.dtype.name}), unit={value.unit})')
+    return value
+
+
+# -------------------------------
+# Graph Definition
+# -------------------------------
+
+
+def _node_flatten(
+    node: Node
+) -> Tuple[Tuple[Tuple[str, Any], ...], Tuple[Type]]:
+    nodes = sorted((key, value) for key, value in vars(node).items() if key != '_trace_state')
+    return nodes, (type(node),)
+
+
+def _node_set_key(
+    node: Node,
+    key: Key,
+    value: Any
+) -> None:
+    if not isinstance(key, str):
+        raise KeyError(f'Invalid key: {key!r}')
+    elif (
+        hasattr(node, key)
+        and isinstance(state := getattr(node, key), State)
+        and isinstance(value, TreefyState)
+    ):
+        state.update_from_ref(value)
+    else:
+        setattr(node, key, value)
+
+
+def _node_pop_key(
+    node: Node,
+    key: Key
+):
+    if not isinstance(key, str):
+        raise KeyError(f'Invalid key: {key!r}')
+    return vars(node).pop(key)
+
+
+def _node_create_empty(
+    static: tuple[Type[G],]
+) -> G:
+    node_type, = static
+    node = object.__new__(node_type)
+    vars(node).update(_trace_state=StateJaxTracer())
+    return node
+
+
+def _node_clear(node: Node):
+    module_state = node._trace_state
+    module_vars = vars(node)
+    module_vars.clear()
+    module_vars['_trace_state'] = module_state
+
+
+class Dict(Node, Mapping[str, A]):
+    """
+    A dictionary node.
+    """
+
+    def __init__(self, *args, **kwargs):
+        for name, value in dict(*args, **kwargs).items():
+            setattr(self, name, value)
+
+    def __getitem__(self, key) -> A:
+        return getattr(self, key)
+
+    def __setitem__(self, key, value):
+        setattr(self, key, value)
+
+    def __getattr__(self, key) -> A:
+        return super().__getattribute__(key)
+
+    def __setattr__(self, key, value):
+        super().__setattr__(key, value)
+
+    def __iter__(self) -> Iterator[str]:
+        return (k for k in vars(self) if k != '_object__state')
+
+    def __len__(self) -> int:
+        return len(vars(self))
+
+
+class List(Node):
+    """
+    A list node.
+    """
+
+    def __init__(self, seq=()):
+        vars(self).update({str(i): item for i, item in enumerate(seq)})
+
+    def __getitem__(self, idx):
+        return getattr(self, str(idx))
+
+    def __setitem__(self, idx, value):
+        setattr(self, str(idx), value)
+
+    def __iter__(self):
+        return iter(vars(self).values())
+
+    def __len__(self):
+        return len(vars(self))
+
+    def __add__(self, other: Sequence[A]) -> List[A]:
+        return List(list(self) + list(other))
+
+    def append(self, value):
+        self[len(vars(self))] = value
+
+    def extend(self, values):
+        for value in values:
+            self.append(value)
+
+
+class Sequential(Node):
+    def __init__(self, *fns: Callable[..., Any]):
+        self.layers = list(fns)
+
+    def __call__(self, *args, **kwargs) -> Any:
+        output: Any = None
+
+        for i, f in enumerate(self.layers):
+            if not callable(f):
+                raise TypeError(f'Sequence[{i}] is not callable: {f}')
+            if i > 0:
+                if isinstance(output, tuple):
+                    args = output
+                    kwargs = {}
+                elif isinstance(output, dict):
+                    args = ()
+                    kwargs = output
+                else:
+                    args = (output,)
+                    kwargs = {}
+            output = f(*args, **kwargs)
+
+        return output

brainstate/graph/_graph_node_test.py

@@ -0,0 +1,75 @@
+# 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
+
+import unittest
+
+import brainstate as bst
+
+
+class TestSequential(unittest.TestCase):
+    def test1(self):
+        s = bst.graph.Sequential(bst.nn.Linear(1, 2),
+                                 bst.nn.Linear(2, 3))
+        graphdef, states = bst.graph.treefy_split(s)
+        print(states)
+        self.assertTrue(len(states.to_flat()) == 2)
+
+
+class TestStateRetrieve(unittest.TestCase):
+    def test_list_of_states_1(self):
+        class Model(bst.graph.Node):
+            def __init__(self):
+                self.a = [1, 2, 3]
+                self.b = [bst.State(1), bst.State(2), bst.State(3)]
+
+        m = Model()
+        graphdef, states = bst.graph.treefy_split(m)
+        print(states.to_flat())
+        self.assertTrue(len(states.to_flat()) == 3)
+
+    def test_list_of_states_2(self):
+        class Model(bst.graph.Node):
+            def __init__(self):
+                self.a = [1, 2, 3]
+                self.b = [bst.State(1), [bst.State(2), bst.State(3)]]
+
+        m = Model()
+        graphdef, states = bst.graph.treefy_split(m)
+        print(states.to_flat())
+        self.assertTrue(len(states.to_flat()) == 3)
+
+    def test_list_of_node_1(self):
+        class Model(bst.graph.Node):
+            def __init__(self):
+                self.a = [1, 2, 3]
+                self.b = [bst.nn.Linear(1, 2), bst.nn.Linear(2, 3)]
+
+        m = Model()
+        graphdef, states = bst.graph.treefy_split(m)
+        print(states.to_flat())
+        self.assertTrue(len(states.to_flat()) == 2)
+
+    def test_list_of_node_2(self):
+        class Model(bst.graph.Node):
+            def __init__(self):
+                self.a = [1, 2, 3]
+                self.b = [bst.nn.Linear(1, 2), [bst.nn.Linear(2, 3)], (bst.nn.Linear(3, 4), bst.nn.Linear(4, 5))]
+
+        m = Model()
+        graphdef, states = bst.graph.treefy_split(m)
+        print(states.to_flat())
+        self.assertTrue(len(states.to_flat()) == 4)