brainstate 0.1.8__py2.py3-none-any.whl → 0.1.9__py2.py3-none-any.whl

brainstate/environ_test.py

@@ -1,62 +1,62 @@
-# 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.
-# ==============================================================================
-
-
-import unittest
-
-import jax.numpy as jnp
-
-import brainstate as bst
-
-
-class TestEnviron(unittest.TestCase):
-    def test_precision(self):
-        with bst.environ.context(precision=64):
-            a = bst.random.randn(1)
-            self.assertEqual(a.dtype, jnp.float64)
-
-        with bst.environ.context(precision=32):
-            a = bst.random.randn(1)
-            self.assertEqual(a.dtype, jnp.float32)
-
-        with bst.environ.context(precision=16):
-            a = bst.random.randn(1)
-            self.assertEqual(a.dtype, jnp.float16)
-
-        with bst.environ.context(precision='bf16'):
-            a = bst.random.randn(1)
-            self.assertEqual(a.dtype, jnp.bfloat16)
-
-    def test_platform(self):
-        with self.assertRaises(ValueError):
-            with bst.environ.context(platform='cpu'):
-                a = bst.random.randn(1)
-                self.assertEqual(a.device(), 'cpu')
-
-    def test_register_default_behavior(self):
-        bst.environ.set(dt=0.1)
-
-        dt_ = 0.1
-
-        def dt_behavior(dt):
-            nonlocal dt_
-            dt_ = dt
-            print(f'dt: {dt}')
-
-        bst.environ.register_default_behavior('dt', dt_behavior)
-
-        with bst.environ.context(dt=0.2):
-            self.assertEqual(dt_, 0.2)
-        self.assertEqual(dt_, 0.1)
+# 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.
+# ==============================================================================
+
+
+import unittest
+
+import jax.numpy as jnp
+
+import brainstate as bst
+
+
+class TestEnviron(unittest.TestCase):
+    def test_precision(self):
+        with bst.environ.context(precision=64):
+            a = bst.random.randn(1)
+            self.assertEqual(a.dtype, jnp.float64)
+
+        with bst.environ.context(precision=32):
+            a = bst.random.randn(1)
+            self.assertEqual(a.dtype, jnp.float32)
+
+        with bst.environ.context(precision=16):
+            a = bst.random.randn(1)
+            self.assertEqual(a.dtype, jnp.float16)
+
+        with bst.environ.context(precision='bf16'):
+            a = bst.random.randn(1)
+            self.assertEqual(a.dtype, jnp.bfloat16)
+
+    def test_platform(self):
+        with self.assertRaises(ValueError):
+            with bst.environ.context(platform='cpu'):
+                a = bst.random.randn(1)
+                self.assertEqual(a.device(), 'cpu')
+
+    def test_register_default_behavior(self):
+        bst.environ.set(dt=0.1)
+
+        dt_ = 0.1
+
+        def dt_behavior(dt):
+            nonlocal dt_
+            dt_ = dt
+            print(f'dt: {dt}')
+
+        bst.environ.register_default_behavior('dt', dt_behavior)
+
+        with bst.environ.context(dt=0.2):
+            self.assertEqual(dt_, 0.2)
+        self.assertEqual(dt_, 0.1)

brainstate/functional/__init__.py

@@ -1,26 +1,27 @@
-# 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 ._activations import *
-from ._activations import __all__ as __activations_all__
-from ._normalization import *
-from ._normalization import __all__ as __others_all__
-from ._others import *
-from ._others import __all__ as __others_all__
-from ._spikes import *
-from ._spikes import __all__ as __spikes_all__
-
-__all__ = __spikes_all__ + __others_all__ + __activations_all__ + __others_all__
+# 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 brainstate.nn._activations import *
+from brainstate.nn._activations import __all__ as act_all
+from brainstate.nn._normalizations import weight_standardization
+from brainstate.nn._others import clip_grad_norm
+
+__all__ = ['weight_standardization', 'clip_grad_norm'] + act_all
+del act_all
+
+if __name__ == '__main__':
+    relu
+    clip_grad_norm
+    weight_standardization

brainstate/graph/__init__.py

@@ -1,29 +1,29 @@
-# 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 ._graph_node import Node, Dict, List, Sequential
-from ._graph_operation import (
-    pop_states, nodes, states, treefy_states, update_states, flatten, unflatten,
-    treefy_split, treefy_merge, iter_leaf, iter_node, clone, graphdef,
-    call, RefMap, GraphDef, NodeRef, NodeDef
-)
-
-__all__ = [
-    'Node', 'Dict', 'List', 'Sequential',
-    'pop_states', 'nodes', 'states', 'treefy_states', 'update_states', 'flatten', 'unflatten',
-    'treefy_split', 'treefy_merge', 'iter_leaf', 'iter_node', 'clone', 'graphdef',
-    'call', 'RefMap', 'GraphDef', 'NodeRef', 'NodeDef',
-]
+# 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 ._graph_node import Node, Dict, List, Sequential
+from ._graph_operation import (
+    pop_states, nodes, states, treefy_states, update_states, flatten, unflatten,
+    treefy_split, treefy_merge, iter_leaf, iter_node, clone, graphdef,
+    call, RefMap, GraphDef, NodeRef, NodeDef
+)
+
+__all__ = [
+    'Node', 'Dict', 'List', 'Sequential',
+    'pop_states', 'nodes', 'states', 'treefy_states', 'update_states', 'flatten', 'unflatten',
+    'treefy_split', 'treefy_merge', 'iter_leaf', 'iter_node', 'clone', 'graphdef',
+    'call', 'RefMap', 'GraphDef', 'NodeRef', 'NodeDef',
+]