[API Compatibility] Improve compat.nn.Linear init and param reset (#76196)

Author: Enigmatisms

python/paddle/compat/nn/__init__.py

@@ -250,20 +250,11 @@ def reset_parameters(self) -> None:
         Resets parameters based on their initialization used in ``__init__``.
         """
 
-        # KaimingUniform initializer should be more flexible: user should be able to specify place
-        expected_place = paddle.base.framework._current_expected_place()
-        original_place = self.weight.place
         nn.init.kaiming_uniform_(self.weight, a=sqrt(5))
-
-        place_mismatch = expected_place != original_place
-        if place_mismatch and in_dynamic_mode():
-            self.weight = self.weight.to(original_place)
         if self.bias is not None:
             # nn.init._calculate_fan_in_and_fan_out(self.weight) for 2D array
             # is equivalent to returning (weight.shape[1], weight.shape[0])
+            # TODO(heqianyue): use _calculate_fan_in_and_fan_out when available
             fan_in = self.weight.shape[1]
             bound = 1 / sqrt(fan_in) if fan_in > 0 else 0
             nn.init.uniform_(self.bias, -bound, bound)
-
-            if place_mismatch and in_dynamic_mode():
-                self.bias = self.bias.to(original_place)

python/paddle/nn/initializer/kaiming.py

@@ -173,12 +173,19 @@ def forward(
                     -limit,
                     limit,
                     self._seed,
-                    _current_expected_place(),
+                    var.place
+                    if var.place._type()
+                    else _current_expected_place(),
                 )
             else:
                 gain = calculate_gain(self._nonlinearity, self._negative_slope)
                 std = gain / math.sqrt(float(fan_in))
-                place = _current_expected_place()
+                # var.place._type() means undefined, happens when initializer is specified in ParamAttr
+                place = (
+                    var.place
+                    if var.place._type()
+                    else _current_expected_place()
+                )
                 out_var = _C_ops.gaussian(
                     out_var.shape, 0.0, std, self._seed, out_dtype, place
                 )

python/paddle/nn/initializer/uniform.py

@@ -126,7 +126,7 @@ def forward(
                 self._low,
                 self._high,
                 self._seed,
-                _current_expected_place(),
+                var.place if var.place._type() else _current_expected_place(),
             )
             if var.dtype == core.VarDesc.VarType.FP16:
                 var_tmp = _C_ops.cast(out_var, var.dtype)

test/legacy_test/test_compat_nn_linear.py

@@ -347,6 +347,18 @@ def test_weight_transpose_behavior(self):
 
         np.testing.assert_allclose(y_pd.numpy(), expected, rtol=1e-5, atol=1e-8)
 
+    def test_reset_parameters(self):
+        if not paddle.base.is_compiled_with_cuda():
+            return
+        devices = ['cpu', None]  # None means the default device
+        for device_ in devices:
+            dummy_tensor = paddle.zeros(1, device=device_)
+            lin = paddle.compat.nn.Linear(4, 8, bias=True, device=device_)
+            expected_device = dummy_tensor.place
+            lin.reset_parameters()
+            self.assertEqual(lin.weight.place, expected_device)
+            self.assertEqual(lin.bias.place, expected_device)
+
     def test_error_handling(self):
         """Test error handling for invalid inputs"""
         # Shape mismatch between input and weight