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

brainstate/init/_random_inits_test.py

@@ -0,0 +1,150 @@
+# 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 -*-
+from __future__ import annotations
+
+import unittest
+
+import brainstate as bst
+
+
+class TestNormalInit(unittest.TestCase):
+
+    def test_normal_init1(self):
+        init = bst.init.Normal()
+        for size in [(100,), (10, 20), (10, 20, 30)]:
+            weights = init(size)
+            assert weights.shape == size
+
+    def test_normal_init2(self):
+        init = bst.init.Normal(scale=0.5)
+        for size in [(100,), (10, 20)]:
+            weights = init(size)
+            assert weights.shape == size
+
+    def test_normal_init3(self):
+        init1 = bst.init.Normal(scale=0.5, seed=10)
+        init2 = bst.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 = bst.init.Normal()
+        for size in [(100,), (10, 20), (10, 20, 30)]:
+            weights = init(size)
+            assert weights.shape == size
+
+    def test_uniform_init2(self):
+        init = bst.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 = bst.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 = bst.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 = bst.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 = bst.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 = bst.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 = bst.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 = bst.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 = bst.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 = bst.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 = bst.init.Orthogonal()
+        for size in [(20, 20), (10, 20, 30)]:
+            weights = init(size)
+            assert weights.shape == size
+
+    def test_orthogonal_init2(self):
+        init = bst.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 = bst.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

@@ -15,6 +15,7 @@
 
 # -*- coding: utf-8 -*-
 
+from __future__ import annotations
 
 import brainunit as u
 
@@ -22,89 +23,85 @@ from brainstate import environ
 from ._base import Initializer, to_size
 
 __all__ = [
-  'ZeroInit',
-  'Constant',
-  'Identity',
+    'ZeroInit',
+    'Constant',
+    'Identity',
 ]
 
 
 class ZeroInit(Initializer):
-  """Zero initializer.
+    """Zero initializer.
 
-  Initialize the weights with zeros.
-  """
+    Initialize the weights with zeros.
+    """
+    __module__ = 'brainstate.init'
 
-  def __init__(self, dtype=None):
-    super(ZeroInit, self).__init__()
-    self.dtype = dtype or environ.dftype()
+    def __init__(self, unit: u.Unit = u.UNITLESS):
+        super(ZeroInit, self).__init__()
+        self.unit = unit
 
-  def __call__(self, shape):
-    shape = to_size(shape)
-    return u.math.zeros(shape, dtype=self.dtype)
-
-  def __repr__(self):
-    return f"{self.__class__.__name__}(dtype={self.dtype})"
+    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.
+    """Constant initializer.
 
-  Parameters
-  ----------
-  value : float, int, bm.ndarray
-    The value to specify.
-  """
+    Initialize the weights with the given values.
 
-  def __init__(self, value=1., dtype=None):
-    super(Constant, self).__init__()
-    self.dtype = dtype or environ.dftype()
-    self.value = value
+    Parameters
+    ----------
+    value : float, int, bm.ndarray
+      The value to specify.
+    """
+    __module__ = 'brainstate.init'
 
-  def __call__(self, shape):
-    shape = to_size(shape)
-    return u.math.full(shape, self.value, dtype=self.dtype)
+    def __init__(self, value=1., ):
+        super(Constant, self).__init__()
+        self.value = value
 
-  def __repr__(self):
-    return f'{self.__class__.__name__}(value={self.value}, dtype={self.dtype})'
+    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).
-  """
-
-  def __init__(self, value=1., dtype=None):
-    super(Identity, self).__init__()
-    self.dtype = dtype or environ.dftype()
-    self.value = value
-
-  def __call__(self, shape):
-    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=self.dtype)
-    r = u.math.fill_diagonal(r, self.value)
-    return r
-
-  def __repr__(self):
-    return f'{self.__class__.__name__}(value={self.value}, dtype={self.dtype})'
+    """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

@@ -0,0 +1,51 @@
+# 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 -*-
+from __future__ import annotations
+
+import unittest
+
+import brainstate as bst
+
+
+class TestZeroInit(unittest.TestCase):
+    def test_zero_init(self):
+        init = bst.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 = bst.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 = bst.init.Identity(value=value)
+                weights = init(size)
+                if len(size) == 1:
+                    assert weights.shape == (size[0], size[0])
+                else:
+                    assert weights.shape == size