【Hackathon 6th No.31】paddle.nn.initializer.Normal support complex number -part (#65105)

NKNaN · web-flow · commit b193ef643185 · 2024-06-24T17:02:34.000+08:00
* update normal initializer

* update test

* fix test
diff --git a/python/paddle/nn/initializer/normal.py b/python/paddle/nn/initializer/normal.py
@@ -33,7 +33,7 @@ class NormalInitializer(Initializer):
     """Implements the Random Normal(Gaussian) distribution initializer
 
     Args:
-        loc (float, optional): mean of the normal distribution. Default is 0.0.
+        loc (float|complex, optional): mean of the normal distribution. Default is 0.0.
         scale (float, optional): standard deviation of the normal distribution. Default is 1.0.
         seed (int, optional): random seed. Default is 0.
 
@@ -49,6 +49,13 @@ def __init__(
         self._mean = loc
         self._std_dev = scale
         self._seed = seed
+        if isinstance(self._mean, complex):
+            if self._mean.real != self._mean.imag:
+                raise ValueError(
+                    "if mean is a complex number, its real part should equal imag part, ",
+                    f"but got real part: {self._mean.real} != imag part: {self._mean.imag}",
+                )
+            self._mean = self._mean.real
 
     def forward(
         self, var: paddle.Tensor, block: pir.Block | None = None
@@ -73,7 +80,14 @@ def forward(
         check_variable_and_dtype(
             var,
             "Out",
-            ["uint16", "float16", "float32", "float64"],
+            [
+                "uint16",
+                "float16",
+                "float32",
+                "float64",
+                "complex64",
+                "complex128",
+            ],
             "guassian_random",
         )
 
@@ -124,7 +138,7 @@ class Normal(NormalInitializer):
     """The Random Normal (Gaussian) distribution initializer.
 
     Args:
-        mean (float, optional): mean of the normal distribution. Default is 0.0.
+        mean (float|complex, optional): mean of the normal distribution. Default is 0.0.
         std (float, optional): standard deviation of the normal distribution. Default is 1.0.
         name(str|None, optional): The default value is None. Normally there is no need for user to set this
             property. For more information, please refer to :ref:`api_guide_Name`. Default: None.
diff --git a/test/deprecated/legacy_test/test_initializer.py b/test/deprecated/legacy_test/test_initializer.py
@@ -425,6 +425,29 @@ def test_normal_initializer(self, dtype="float32"):
         self.assertEqual(init_op.attr('seed'), 123)
         return block
 
+    def test_normal_initializer_complex(self, dtype="complex64"):
+        """Test normal initializer with complex dtype"""
+        program = framework.Program()
+        block = program.global_block()
+        for _ in range(2):
+            block.create_parameter(
+                dtype=dtype,
+                shape=[5, 10],
+                lod_level=0,
+                name="param",
+                initializer=paddle.nn.initializer.NormalInitializer(
+                    2.2 + 2.2j, 1.9, 123
+                ),
+            )
+        num_ops = 1
+        self.assertEqual(len(block.ops), num_ops)
+        init_op = block.ops[0]
+        self.assertEqual(init_op.type, 'gaussian_random')
+        self.assertAlmostEqual(init_op.attr('mean'), 2.2, delta=DELTA)
+        self.assertAlmostEqual(init_op.attr('std'), 1.9, delta=DELTA)
+        self.assertEqual(init_op.attr('seed'), 123)
+        return block
+
     def test_normal_initializer_fp16(self):
         """Test normal initializer with float16"""
         self.test_normal_initializer("float16")
@@ -433,6 +456,14 @@ def test_normal_initializer_bf16(self):
         """Test normal initializer with bfloat16"""
         self.test_normal_initializer("uint16")
 
+    def test_normal_initializer_complex64(self):
+        """Test normal initializer with complex64"""
+        self.test_normal_initializer_complex("complex64")
+
+    def test_normal_initializer_complex128(self):
+        """Test normal initializer with complex128"""
+        self.test_normal_initializer_complex("complex128")
+
 
 class TestXavierInitializer(unittest.TestCase):
     def test_uniform_xavier_initializer(self):
diff --git a/test/deprecated/legacy_test/test_initializer_nn.py b/test/deprecated/legacy_test/test_initializer_nn.py
@@ -409,6 +409,28 @@ def test_normal_initializer(self, dtype="float32"):
 
         return block
 
+    def test_normal_initializer_complex(self, dtype="complex64"):
+        """Test normal initializer with complex dtype"""
+        with static_guard():
+            program = framework.Program()
+            block = program.global_block()
+            for _ in range(2):
+                block.create_parameter(
+                    dtype=dtype,
+                    shape=[5, 10],
+                    lod_level=0,
+                    name="param",
+                    initializer=initializer.Normal(2.3 + 2.3j, 1.9),
+                )
+            num_ops = 1
+            self.assertEqual(len(block.ops), num_ops)
+            init_op = block.ops[0]
+            self.assertEqual(init_op.type, 'gaussian_random')
+            self.assertAlmostEqual(init_op.attr('mean'), 2.3, delta=DELTA)
+            self.assertAlmostEqual(init_op.attr('std'), 1.9, delta=DELTA)
+
+        return block
+
     def test_normal_initializer_fp16(self):
         """Test normal initializer with float16"""
         block = self.test_normal_initializer("float16")
@@ -417,6 +439,14 @@ def test_normal_initializer_bf16(self):
         """Test normal initializer with bfloat16"""
         block = self.test_normal_initializer("uint16")  # bfloat16
 
+    def test_normal_initializer_complex64(self):
+        """Test normal initializer with complex64"""
+        block = self.test_normal_initializer_complex("complex64")
+
+    def test_normal_initializer_complex128(self):
+        """Test normal initializer with complex128"""
+        block = self.test_normal_initializer_complex("complex128")
+
     def test_normal_initializer_dygraph(self):
         """Test normal initializer in dygraph model."""
         paddle.disable_static()
@@ -427,6 +457,10 @@ def test_normal_initializer_dygraph(self):
         )
         linear = paddle.nn.Linear(2, 2, weight_attr=weight_attr)
 
+    def test_normal_initializer_complex_error(self):
+        with self.assertRaises(ValueError):
+            paddle.nn.initializer.Normal(mean=0.1 + 0.2j, std=2.0)
+
 
 class TestTruncatedNormal(unittest.TestCase):
     def test_truncated_normal_initializer_default_value(self):
