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

brainstate/util/_dict_test.py

@@ -0,0 +1,160 @@
+# 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 jax
+from absl.testing import absltest
+
+import brainstate as bst
+
+
+class TestNestedMapping(absltest.TestCase):
+    def test_create_state(self):
+        state = bst.util.NestedDict({'a': bst.ParamState(1), 'b': {'c': bst.ParamState(2)}})
+
+        assert state['a'].value == 1
+        assert state['b']['c'].value == 2
+
+    def test_get_attr(self):
+        state = bst.util.NestedDict({'a': bst.ParamState(1), 'b': {'c': bst.ParamState(2)}})
+
+        assert state.a.value == 1
+        assert state.b['c'].value == 2
+
+    def test_set_attr(self):
+        state = bst.util.NestedDict({'a': bst.ParamState(1), 'b': {'c': bst.ParamState(2)}})
+
+        state.a.value = 3
+        state.b['c'].value = 4
+
+        assert state['a'].value == 3
+        assert state['b']['c'].value == 4
+
+    def test_set_attr_variables(self):
+        state = bst.util.NestedDict({'a': bst.ParamState(1), 'b': {'c': bst.ParamState(2)}})
+
+        state.a.value = 3
+        state.b['c'].value = 4
+
+        assert isinstance(state.a, bst.ParamState)
+        assert state.a.value == 3
+        assert isinstance(state.b['c'], bst.ParamState)
+        assert state.b['c'].value == 4
+
+    def test_add_nested_attr(self):
+        state = bst.util.NestedDict({'a': bst.ParamState(1), 'b': {'c': bst.ParamState(2)}})
+        state.b['d'] = bst.ParamState(5)
+
+        assert state['b']['d'].value == 5
+
+    def test_delete_nested_attr(self):
+        state = bst.util.NestedDict({'a': bst.ParamState(1), 'b': {'c': bst.ParamState(2)}})
+        del state['b']['c']
+
+        assert 'c' not in state['b']
+
+    def test_integer_access(self):
+        class Foo(bst.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layers = [bst.nn.Linear(1, 2), bst.nn.Linear(2, 3)]
+
+        module = Foo()
+        state_refs = bst.graph.treefy_states(module)
+
+        assert module.layers[0].weight.value['weight'].shape == (1, 2)
+        assert state_refs.layers[0]['weight'].value['weight'].shape == (1, 2)
+        assert module.layers[1].weight.value['weight'].shape == (2, 3)
+        assert state_refs.layers[1]['weight'].value['weight'].shape == (2, 3)
+
+    def test_pure_dict(self):
+        module = bst.nn.Linear(4, 5)
+        state_map = bst.graph.treefy_states(module)
+        pure_dict = state_map.to_pure_dict()
+        assert isinstance(pure_dict, dict)
+        assert isinstance(pure_dict['weight'].value['weight'], jax.Array)
+        assert isinstance(pure_dict['weight'].value['bias'], jax.Array)
+
+
+class TestSplit(unittest.TestCase):
+    def test_split(self):
+        class Model(bst.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.batchnorm = bst.nn.BatchNorm1d([10, 3])
+                self.linear = bst.nn.Linear([10, 3], [10, 4])
+
+            def __call__(self, x):
+                return self.linear(self.batchnorm(x))
+
+        with bst.environ.context(fit=True):
+            model = Model()
+            x = bst.random.randn(1, 10, 3)
+            y = model(x)
+            self.assertEqual(y.shape, (1, 10, 4))
+
+        state_map = bst.graph.treefy_states(model)
+
+        with self.assertRaises(ValueError):
+            params, others = state_map.split(bst.ParamState)
+
+        params, others = state_map.split(bst.ParamState, ...)
+        print()
+        print(params)
+        print(others)
+
+        self.assertTrue(len(params.to_flat()) == 2)
+        self.assertTrue(len(others.to_flat()) == 2)
+
+
+class TestStateMap2(unittest.TestCase):
+    def test1(self):
+        class Model(bst.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.batchnorm = bst.nn.BatchNorm1d([10, 3])
+                self.linear = bst.nn.Linear([10, 3], [10, 4])
+
+            def __call__(self, x):
+                return self.linear(self.batchnorm(x))
+
+        with bst.environ.context(fit=True):
+            model = Model()
+            state_map = bst.graph.treefy_states(model).to_flat()
+            state_map = bst.util.NestedDict(state_map)
+
+
+class TestFlattedMapping(unittest.TestCase):
+    def test1(self):
+        class Model(bst.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.batchnorm = bst.nn.BatchNorm1d([10, 3])
+                self.linear = bst.nn.Linear([10, 3], [10, 4])
+
+            def __call__(self, x):
+                return self.linear(self.batchnorm(x))
+
+        model = Model()
+        # print(model.states())
+        # print(bst.graph.states(model))
+        self.assertTrue(model.states() == bst.graph.states(model))
+
+        print(model.nodes())
+        # print(bst.graph.nodes(model))
+        self.assertTrue(model.nodes() == bst.graph.nodes(model))

brainstate/util/_error.py

@@ -0,0 +1,28 @@
+# 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.
+# ==============================================================================
+
+
+__all__ = [
+    'BrainStateError',
+    'TraceContextError',
+]
+
+
+class BrainStateError(Exception):
+    pass
+
+
+class TraceContextError(BrainStateError):
+    pass

brainstate/util/_filter.py

@@ -0,0 +1,178 @@
+# 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
+
+import builtins
+import dataclasses
+import typing
+from typing import TYPE_CHECKING
+
+from brainstate.typing import Filter, PathParts, Predicate, Key
+
+if TYPE_CHECKING:
+    ellipsis = builtins.ellipsis
+else:
+    ellipsis = typing.Any
+
+__all__ = [
+    'to_predicate',
+]
+
+
+def to_predicate(the_filter: Filter) -> Predicate:
+    """
+    Converts a Filter to a predicate function.
+    """
+
+    if isinstance(the_filter, str):
+        return WithTagFilter(the_filter)
+    elif isinstance(the_filter, type):
+        return OfTypeFilter(the_filter)
+    elif isinstance(the_filter, bool):
+        if the_filter:
+            return EverythingFilter()
+        else:
+            return NothingFilter()
+    elif the_filter is Ellipsis:
+        return EverythingFilter()
+    elif the_filter is None:
+        return NothingFilter()
+    elif callable(the_filter):
+        return the_filter
+    elif isinstance(the_filter, (list, tuple)):
+        return AnyFilter(*the_filter)
+    else:
+        raise TypeError(f'Invalid collection filter: {the_filter:!r}. ')
+
+
+@dataclasses.dataclass(frozen=True)
+class WithTagFilter:
+    tag: str
+
+    def __call__(self, path: PathParts, x: typing.Any):
+        return hasattr(x, 'tag') and x.tag == self.tag
+
+    def __repr__(self):
+        return f'WithTag({self.tag!r})'
+
+
+@dataclasses.dataclass(frozen=True)
+class PathContainsFilter:
+    key: Key
+
+    def __call__(self, path: PathParts, x: typing.Any):
+        return self.key in path
+
+    def __repr__(self):
+        return f'PathContains({self.key!r})'
+
+
+@dataclasses.dataclass(frozen=True)
+class OfTypeFilter:
+    type: type
+
+    def __call__(self, path: PathParts, x: typing.Any):
+        return isinstance(x, self.type) or (
+            hasattr(x, 'type') and issubclass(x.type, self.type)
+        )
+
+    def __repr__(self):
+        return f'OfType({self.type!r})'
+
+
+class AnyFilter:
+    def __init__(self, *filters: Filter):
+        self.predicates = tuple(
+            to_predicate(collection_filter) for collection_filter in filters
+        )
+
+    def __call__(self, path: PathParts, x: typing.Any):
+        return any(predicate(path, x) for predicate in self.predicates)
+
+    def __repr__(self):
+        return f'Any({", ".join(map(repr, self.predicates))})'
+
+    def __eq__(self, other):
+        return isinstance(other, AnyFilter) and self.predicates == other.predicates
+
+    def __hash__(self):
+        return hash(self.predicates)
+
+
+class AllFilter:
+    def __init__(self, *filters: Filter):
+        self.predicates = tuple(
+            to_predicate(collection_filter) for collection_filter in filters
+        )
+
+    def __call__(self, path: PathParts, x: typing.Any):
+        return all(predicate(path, x) for predicate in self.predicates)
+
+    def __repr__(self):
+        return f'All({", ".join(map(repr, self.predicates))})'
+
+    def __eq__(self, other):
+        return isinstance(other, AllFilter) and self.predicates == other.predicates
+
+    def __hash__(self):
+        return hash(self.predicates)
+
+
+class NotFilter:
+    def __init__(self, collection_filter: Filter, /):
+        self.predicate = to_predicate(collection_filter)
+
+    def __call__(self, path: PathParts, x: typing.Any):
+        return not self.predicate(path, x)
+
+    def __repr__(self):
+        return f'Not({self.predicate!r})'
+
+    def __eq__(self, other):
+        return isinstance(other, NotFilter) and self.predicate == other.predicate
+
+    def __hash__(self):
+        return hash(self.predicate)
+
+
+class EverythingFilter:
+    def __call__(self, path: PathParts, x: typing.Any):
+        return True
+
+    def __repr__(self):
+        return 'Everything()'
+
+    def __eq__(self, other):
+        return isinstance(other, EverythingFilter)
+
+    def __hash__(self):
+        return hash(EverythingFilter)
+
+
+class NothingFilter:
+    def __call__(self, path: PathParts, x: typing.Any):
+        return False
+
+    def __repr__(self):
+        return 'Nothing()'
+
+    def __eq__(self, other):
+        return isinstance(other, NothingFilter)
+
+    def __hash__(self):
+        return hash(NothingFilter)