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

brainstate/init/_random_inits_test.py

@@ -1,149 +1,149 @@
-# 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.
-# ==============================================================================
-
-# -*- coding: utf-8 -*-
-
-import unittest
-
-import brainstate
-
-
-class TestNormalInit(unittest.TestCase):
-
-    def test_normal_init1(self):
-        init = brainstate.init.Normal()
-        for size in [(100,), (10, 20), (10, 20, 30)]:
-            weights = init(size)
-            assert weights.shape == size
-
-    def test_normal_init2(self):
-        init = brainstate.init.Normal(scale=0.5)
-        for size in [(100,), (10, 20)]:
-            weights = init(size)
-            assert weights.shape == size
-
-    def test_normal_init3(self):
-        init1 = brainstate.init.Normal(scale=0.5, seed=10)
-        init2 = brainstate.init.Normal(scale=0.5, seed=10)
-        size = (10,)
-        weights1 = init1(size)
-        weights2 = init2(size)
-        assert weights1.shape == size
-        assert (weights1 == weights2).all()
-
-
-class TestUniformInit(unittest.TestCase):
-    def test_uniform_init1(self):
-        init = brainstate.init.Normal()
-        for size in [(100,), (10, 20), (10, 20, 30)]:
-            weights = init(size)
-            assert weights.shape == size
-
-    def test_uniform_init2(self):
-        init = brainstate.init.Uniform(min_val=10, max_val=20)
-        for size in [(100,), (10, 20)]:
-            weights = init(size)
-            assert weights.shape == size
-
-
-class TestVarianceScaling(unittest.TestCase):
-    def test_var_scaling1(self):
-        init = brainstate.init.VarianceScaling(scale=1., mode='fan_in', distribution='truncated_normal')
-        for size in [(10, 20), (10, 20, 30)]:
-            weights = init(size)
-            assert weights.shape == size
-
-    def test_var_scaling2(self):
-        init = brainstate.init.VarianceScaling(scale=2, mode='fan_out', distribution='normal')
-        for size in [(10, 20), (10, 20, 30)]:
-            weights = init(size)
-            assert weights.shape == size
-
-    def test_var_scaling3(self):
-        init = brainstate.init.VarianceScaling(scale=2 / 4, mode='fan_avg', in_axis=0, out_axis=1,
-                                               distribution='uniform')
-        for size in [(10, 20), (10, 20, 30)]:
-            weights = init(size)
-            assert weights.shape == size
-
-
-class TestKaimingUniformUnit(unittest.TestCase):
-    def test_kaiming_uniform_init(self):
-        init = brainstate.init.KaimingUniform()
-        for size in [(10, 20), (10, 20, 30)]:
-            weights = init(size)
-            assert weights.shape == size
-
-
-class TestKaimingNormalUnit(unittest.TestCase):
-    def test_kaiming_normal_init(self):
-        init = brainstate.init.KaimingNormal()
-        for size in [(10, 20), (10, 20, 30)]:
-            weights = init(size)
-            assert weights.shape == size
-
-
-class TestXavierUniformUnit(unittest.TestCase):
-    def test_xavier_uniform_init(self):
-        init = brainstate.init.XavierUniform()
-        for size in [(10, 20), (10, 20, 30)]:
-            weights = init(size)
-            assert weights.shape == size
-
-
-class TestXavierNormalUnit(unittest.TestCase):
-    def test_xavier_normal_init(self):
-        init = brainstate.init.XavierNormal()
-        for size in [(10, 20), (10, 20, 30)]:
-            weights = init(size)
-            assert weights.shape == size
-
-
-class TestLecunUniformUnit(unittest.TestCase):
-    def test_lecun_uniform_init(self):
-        init = brainstate.init.LecunUniform()
-        for size in [(10, 20), (10, 20, 30)]:
-            weights = init(size)
-            assert weights.shape == size
-
-
-class TestLecunNormalUnit(unittest.TestCase):
-    def test_lecun_normal_init(self):
-        init = brainstate.init.LecunNormal()
-        for size in [(10, 20), (10, 20, 30)]:
-            weights = init(size)
-            assert weights.shape == size
-
-
-class TestOrthogonalUnit(unittest.TestCase):
-    def test_orthogonal_init1(self):
-        init = brainstate.init.Orthogonal()
-        for size in [(20, 20), (10, 20, 30)]:
-            weights = init(size)
-            assert weights.shape == size
-
-    def test_orthogonal_init2(self):
-        init = brainstate.init.Orthogonal(scale=2., axis=0)
-        for size in [(10, 20), (10, 20, 30)]:
-            weights = init(size)
-            assert weights.shape == size
-
-
-class TestDeltaOrthogonalUnit(unittest.TestCase):
-    def test_delta_orthogonal_init1(self):
-        init = brainstate.init.DeltaOrthogonal()
-        for size in [(20, 20, 20), (10, 20, 30, 40), (50, 40, 30, 20, 20)]:
-            weights = init(size)
-            assert weights.shape == size
+# 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.
+# ==============================================================================
+
+# -*- coding: utf-8 -*-
+
+import unittest
+
+import brainstate
+
+
+class TestNormalInit(unittest.TestCase):
+
+    def test_normal_init1(self):
+        init = brainstate.init.Normal()
+        for size in [(100,), (10, 20), (10, 20, 30)]:
+            weights = init(size)
+            assert weights.shape == size
+
+    def test_normal_init2(self):
+        init = brainstate.init.Normal(scale=0.5)
+        for size in [(100,), (10, 20)]:
+            weights = init(size)
+            assert weights.shape == size
+
+    def test_normal_init3(self):
+        init1 = brainstate.init.Normal(scale=0.5, seed=10)
+        init2 = brainstate.init.Normal(scale=0.5, seed=10)
+        size = (10,)
+        weights1 = init1(size)
+        weights2 = init2(size)
+        assert weights1.shape == size
+        assert (weights1 == weights2).all()
+
+
+class TestUniformInit(unittest.TestCase):
+    def test_uniform_init1(self):
+        init = brainstate.init.Normal()
+        for size in [(100,), (10, 20), (10, 20, 30)]:
+            weights = init(size)
+            assert weights.shape == size
+
+    def test_uniform_init2(self):
+        init = brainstate.init.Uniform(min_val=10, max_val=20)
+        for size in [(100,), (10, 20)]:
+            weights = init(size)
+            assert weights.shape == size
+
+
+class TestVarianceScaling(unittest.TestCase):
+    def test_var_scaling1(self):
+        init = brainstate.init.VarianceScaling(scale=1., mode='fan_in', distribution='truncated_normal')
+        for size in [(10, 20), (10, 20, 30)]:
+            weights = init(size)
+            assert weights.shape == size
+
+    def test_var_scaling2(self):
+        init = brainstate.init.VarianceScaling(scale=2, mode='fan_out', distribution='normal')
+        for size in [(10, 20), (10, 20, 30)]:
+            weights = init(size)
+            assert weights.shape == size
+
+    def test_var_scaling3(self):
+        init = brainstate.init.VarianceScaling(scale=2 / 4, mode='fan_avg', in_axis=0, out_axis=1,
+                                               distribution='uniform')
+        for size in [(10, 20), (10, 20, 30)]:
+            weights = init(size)
+            assert weights.shape == size
+
+
+class TestKaimingUniformUnit(unittest.TestCase):
+    def test_kaiming_uniform_init(self):
+        init = brainstate.init.KaimingUniform()
+        for size in [(10, 20), (10, 20, 30)]:
+            weights = init(size)
+            assert weights.shape == size
+
+
+class TestKaimingNormalUnit(unittest.TestCase):
+    def test_kaiming_normal_init(self):
+        init = brainstate.init.KaimingNormal()
+        for size in [(10, 20), (10, 20, 30)]:
+            weights = init(size)
+            assert weights.shape == size
+
+
+class TestXavierUniformUnit(unittest.TestCase):
+    def test_xavier_uniform_init(self):
+        init = brainstate.init.XavierUniform()
+        for size in [(10, 20), (10, 20, 30)]:
+            weights = init(size)
+            assert weights.shape == size
+
+
+class TestXavierNormalUnit(unittest.TestCase):
+    def test_xavier_normal_init(self):
+        init = brainstate.init.XavierNormal()
+        for size in [(10, 20), (10, 20, 30)]:
+            weights = init(size)
+            assert weights.shape == size
+
+
+class TestLecunUniformUnit(unittest.TestCase):
+    def test_lecun_uniform_init(self):
+        init = brainstate.init.LecunUniform()
+        for size in [(10, 20), (10, 20, 30)]:
+            weights = init(size)
+            assert weights.shape == size
+
+
+class TestLecunNormalUnit(unittest.TestCase):
+    def test_lecun_normal_init(self):
+        init = brainstate.init.LecunNormal()
+        for size in [(10, 20), (10, 20, 30)]:
+            weights = init(size)
+            assert weights.shape == size
+
+
+class TestOrthogonalUnit(unittest.TestCase):
+    def test_orthogonal_init1(self):
+        init = brainstate.init.Orthogonal()
+        for size in [(20, 20), (10, 20, 30)]:
+            weights = init(size)
+            assert weights.shape == size
+
+    def test_orthogonal_init2(self):
+        init = brainstate.init.Orthogonal(scale=2., axis=0)
+        for size in [(10, 20), (10, 20, 30)]:
+            weights = init(size)
+            assert weights.shape == size
+
+
+class TestDeltaOrthogonalUnit(unittest.TestCase):
+    def test_delta_orthogonal_init1(self):
+        init = brainstate.init.DeltaOrthogonal()
+        for size in [(20, 20, 20), (10, 20, 30, 40), (50, 40, 30, 20, 20)]:
+            weights = init(size)
+            assert weights.shape == size

brainstate/init/_regular_inits.py

@@ -1,105 +1,105 @@
-# 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.
-# ==============================================================================
-
-# -*- coding: utf-8 -*-
-
-import brainunit as u
-
-from brainstate import environ
-from ._base import Initializer, to_size
-
-__all__ = [
-    'ZeroInit',
-    'Constant',
-    'Identity',
-]
-
-
-class ZeroInit(Initializer):
-    """Zero initializer.
-
-    Initialize the weights with zeros.
-    """
-    __module__ = 'brainstate.init'
-
-    def __init__(self, unit: u.Unit = u.UNITLESS):
-        super(ZeroInit, self).__init__()
-        self.unit = unit
-
-    def __call__(self, shape, dtype=None):
-        dtype = dtype or environ.dftype()
-        shape = to_size(shape)
-        return u.maybe_decimal(u.math.zeros(shape, dtype=dtype, unit=self.unit))
-
-
-class Constant(Initializer):
-    """Constant initializer.
-
-    Initialize the weights with the given values.
-
-    Parameters
-    ----------
-    value : float, int, bm.ndarray
-      The value to specify.
-    """
-    __module__ = 'brainstate.init'
-
-    def __init__(self, value=1., ):
-        super(Constant, self).__init__()
-        self.value = value
-
-    def __call__(self, shape, dtype=None):
-        dtype = dtype or environ.dftype()
-        shape = to_size(shape)
-        return u.maybe_decimal(u.math.full(shape, self.value, dtype=dtype))
-
-
-class Identity(Initializer):
-    """Returns the identity matrix.
-
-    This initializer was proposed in (Le, et al., 2015) [1]_.
-
-    Parameters
-    ----------
-    value : float
-      The optional scaling factor.
-
-    Returns
-    -------
-    shape: tuple of int
-      The weight shape/size.
-
-    References
-    ----------
-    .. [1] Le, Quoc V., Navdeep Jaitly, and Geoffrey E. Hinton. "A simple way to
-           initialize recurrent networks of rectified linear units." arXiv preprint
-           arXiv:1504.00941 (2015).
-    """
-    __module__ = 'brainstate.init'
-
-    def __init__(self, value=1., unit: u.Unit = u.UNITLESS):
-        super(Identity, self).__init__()
-        self.value = value
-        self.unit = unit
-
-    def __call__(self, shape, dtype=None):
-        dtype = dtype or environ.dftype()
-        shape = to_size(shape)
-        if isinstance(shape, (tuple, list)):
-            if len(shape) > 2:
-                raise ValueError(f'Only support initialize 2D weights for {self.__class__.__name__}.')
-        r = u.math.eye(*shape, dtype=dtype)
-        r = u.math.fill_diagonal(r, self.value)
-        return u.maybe_decimal(u.Quantity(r, unit=self.unit))
+# 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.
+# ==============================================================================
+
+# -*- coding: utf-8 -*-
+
+import brainunit as u
+
+from brainstate import environ
+from ._base import Initializer, to_size
+
+__all__ = [
+    'ZeroInit',
+    'Constant',
+    'Identity',
+]
+
+
+class ZeroInit(Initializer):
+    """Zero initializer.
+
+    Initialize the weights with zeros.
+    """
+    __module__ = 'brainstate.init'
+
+    def __init__(self, unit: u.Unit = u.UNITLESS):
+        super(ZeroInit, self).__init__()
+        self.unit = unit
+
+    def __call__(self, shape, dtype=None):
+        dtype = dtype or environ.dftype()
+        shape = to_size(shape)
+        return u.maybe_decimal(u.math.zeros(shape, dtype=dtype, unit=self.unit))
+
+
+class Constant(Initializer):
+    """Constant initializer.
+
+    Initialize the weights with the given values.
+
+    Parameters
+    ----------
+    value : float, int, bm.ndarray
+      The value to specify.
+    """
+    __module__ = 'brainstate.init'
+
+    def __init__(self, value=1., ):
+        super(Constant, self).__init__()
+        self.value = value
+
+    def __call__(self, shape, dtype=None):
+        dtype = dtype or environ.dftype()
+        shape = to_size(shape)
+        return u.maybe_decimal(u.math.full(shape, self.value, dtype=dtype))
+
+
+class Identity(Initializer):
+    """Returns the identity matrix.
+
+    This initializer was proposed in (Le, et al., 2015) [1]_.
+
+    Parameters
+    ----------
+    value : float
+      The optional scaling factor.
+
+    Returns
+    -------
+    shape: tuple of int
+      The weight shape/size.
+
+    References
+    ----------
+    .. [1] Le, Quoc V., Navdeep Jaitly, and Geoffrey E. Hinton. "A simple way to
+           initialize recurrent networks of rectified linear units." arXiv preprint
+           arXiv:1504.00941 (2015).
+    """
+    __module__ = 'brainstate.init'
+
+    def __init__(self, value=1., unit: u.Unit = u.UNITLESS):
+        super(Identity, self).__init__()
+        self.value = value
+        self.unit = unit
+
+    def __call__(self, shape, dtype=None):
+        dtype = dtype or environ.dftype()
+        shape = to_size(shape)
+        if isinstance(shape, (tuple, list)):
+            if len(shape) > 2:
+                raise ValueError(f'Only support initialize 2D weights for {self.__class__.__name__}.')
+        r = u.math.eye(*shape, dtype=dtype)
+        r = u.math.fill_diagonal(r, self.value)
+        return u.maybe_decimal(u.Quantity(r, unit=self.unit))

brainstate/init/_regular_inits_test.py

@@ -1,50 +1,50 @@
-# 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.
-# ==============================================================================
-
-# -*- coding: utf-8 -*-
-
-import unittest
-
-import brainstate
-
-
-class TestZeroInit(unittest.TestCase):
-    def test_zero_init(self):
-        init = brainstate.init.ZeroInit()
-        for size in [(100,), (10, 20), (10, 20, 30)]:
-            weights = init(size)
-            assert weights.shape == size
-
-
-class TestOneInit(unittest.TestCase):
-    def test_one_init(self):
-        for size in [(100,), (10, 20), (10, 20, 30)]:
-            for value in [0., 1., -1.]:
-                init = brainstate.init.Constant(value=value)
-                weights = init(size)
-                assert weights.shape == size
-                assert (weights == value).all()
-
-
-class TestIdentityInit(unittest.TestCase):
-    def test_identity_init(self):
-        for size in [(100,), (10, 20)]:
-            for value in [0., 1., -1.]:
-                init = brainstate.init.Identity(value=value)
-                weights = init(size)
-                if len(size) == 1:
-                    assert weights.shape == (size[0], size[0])
-                else:
-                    assert weights.shape == size
+# 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.
+# ==============================================================================
+
+# -*- coding: utf-8 -*-
+
+import unittest
+
+import brainstate
+
+
+class TestZeroInit(unittest.TestCase):
+    def test_zero_init(self):
+        init = brainstate.init.ZeroInit()
+        for size in [(100,), (10, 20), (10, 20, 30)]:
+            weights = init(size)
+            assert weights.shape == size
+
+
+class TestOneInit(unittest.TestCase):
+    def test_one_init(self):
+        for size in [(100,), (10, 20), (10, 20, 30)]:
+            for value in [0., 1., -1.]:
+                init = brainstate.init.Constant(value=value)
+                weights = init(size)
+                assert weights.shape == size
+                assert (weights == value).all()
+
+
+class TestIdentityInit(unittest.TestCase):
+    def test_identity_init(self):
+        for size in [(100,), (10, 20)]:
+            for value in [0., 1., -1.]:
+                init = brainstate.init.Identity(value=value)
+                weights = init(size)
+                if len(size) == 1:
+                    assert weights.shape == (size[0], size[0])
+                else:
+                    assert weights.shape == size