【PIR API adaptor No.238、239、240、241】 Migrate nn.initializer.XavierIni…

…tializer, nn.initializer.MSRAInitializer into pir (PaddlePaddle#59419)

python/paddle/nn/initializer/kaiming.py

@@ -15,10 +15,15 @@
 # TODO: define the initializers of Kaiming functions in neural network
 import math
 
+import paddle
 from paddle import _C_ops
 
 from ...base import core, framework, unique_name
-from ...base.framework import _current_expected_place, in_dygraph_mode
+from ...base.framework import (
+    _current_expected_place,
+    in_dygraph_mode,
+    in_pir_mode,
+)
 from .initializer import Initializer, calculate_gain
 
 __all__ = []
@@ -87,9 +92,10 @@ def forward(self, var, block=None):
             The initialization op.
         """
         block = self._check_block(block)
-
-        assert isinstance(var, framework.Variable)
-        assert isinstance(block, framework.Block)
+        assert isinstance(
+            var, (framework.Variable, paddle.pir.core.ParameterMeta)
+        )
+        assert isinstance(block, (framework.Block, paddle.pir.Block))
         f_in, f_out = self._compute_fans(var)
 
         # If fan_in is passed, use it
@@ -112,6 +118,12 @@ def forward(self, var, block=None):
                 type=core.VarDesc.VarType.LOD_TENSOR,
                 persistable=False,
             )
+        elif (
+            var.dtype in (core.DataType.FLOAT16, core.DataType.BFLOAT16)
+            and not self._uniform
+        ):
+            out_dtype = core.DataType.FLOAT32
+            out_var = var
         else:
             out_dtype = var.dtype
             out_var = var
@@ -144,6 +156,33 @@ def forward(self, var, block=None):
             else:
                 out_var._share_underline_tensor_to(var)
             return None
+        elif in_pir_mode():
+            if self._uniform:
+                gain = calculate_gain(self._nonlinearity, self._negative_slope)
+                limit = gain * math.sqrt(3.0 / float(fan_in))
+                out_var = _C_ops.uniform(
+                    var.shape,
+                    out_dtype,
+                    -limit,
+                    limit,
+                    self._seed,
+                    _current_expected_place(),
+                )
+            else:
+                gain = calculate_gain(self._nonlinearity, self._negative_slope)
+                std = gain / math.sqrt(float(fan_in))
+                place = _current_expected_place()
+                out_var = _C_ops.gaussian(
+                    out_var.shape, 0.0, std, self._seed, out_dtype, place
+                )
+
+            if (
+                var.dtype in (core.DataType.FLOAT16, core.DataType.BFLOAT16)
+                and not self._uniform
+            ):
+                return _C_ops.cast(out_var, var.dtype)
+
+            return out_var
         else:
             if self._uniform:
                 gain = calculate_gain(self._nonlinearity, self._negative_slope)

python/paddle/nn/initializer/xavier.py

@@ -119,6 +119,12 @@ def forward(self, var, block=None):
                 type=core.VarDesc.VarType.LOD_TENSOR,
                 persistable=False,
             )
+        elif (
+            var.dtype in (core.DataType.FLOAT16, core.DataType.BFLOAT16)
+            and not self._uniform
+        ):
+            out_dtype = core.DataType.FLOAT32
+            out_var = var
         else:
             out_dtype = var.dtype
             out_var = var
@@ -153,14 +159,32 @@ def forward(self, var, block=None):
         elif in_pir_mode():
             if self._uniform:
                 limit = math.sqrt(6.0 / float(fan_in + fan_out))
-                return paddle._pir_ops.uniform(
-                    var.shape,
-                    var.dtype,
+                out_var = paddle._pir_ops.uniform(
+                    out_var.shape,
+                    out_dtype,
                     -limit,
                     limit,
                     self._seed,
                     _current_expected_place(),
                 )
+            else:
+                std = math.sqrt(2.0 / float(fan_in + fan_out))
+                out_var = _C_ops.gaussian(
+                    out_var.shape,
+                    0.0,
+                    std,
+                    self._seed,
+                    out_dtype,
+                    _current_expected_place(),
+                )
+
+            if (
+                var.dtype in (core.DataType.FLOAT16, core.DataType.BFLOAT16)
+                and not self._uniform
+            ):
+                return _C_ops.cast(out_var, var.dtype)
+
+            return out_var
         else:
             if self._uniform:
                 limit = math.sqrt(6.0 / float(fan_in + fan_out))