Fix matrix_rank_atol_rtol kernel when rtol = None 易用性提升 (#68108)

NKNaN · web-flow · commit 5d2fd0455fda · 2024-09-14T14:09:25.000+08:00
diff --git a/paddle/phi/kernels/cpu/matrix_rank_tol_kernel.cc b/paddle/phi/kernels/cpu/matrix_rank_tol_kernel.cc
@@ -227,23 +227,24 @@ void MatrixRankAtolRtolKernel(const Context& dev_ctx,
     // when `rtol` is specified to be None in py api
     // use rtol=eps*max(m, n) only if `atol` is passed with value 0.0, else use
     // rtol=0.0
-    DenseTensor zero_tensor;
-    zero_tensor = phi::FullLike<T, Context>(dev_ctx, atol, static_cast<T>(0.0));
-
     T rtol_T = std::numeric_limits<T>::epsilon() * std::max(rows, cols);
     DenseTensor default_rtol_tensor;
     default_rtol_tensor =
         phi::Full<T, Context>(dev_ctx, {1}, static_cast<T>(rtol_T));
     default_rtol_tensor =
         phi::Multiply<T>(dev_ctx, default_rtol_tensor, max_eigenvalue_tensor);
 
+    DenseTensor zero_tensor;
+    zero_tensor = phi::FullLike<T, Context>(
+        dev_ctx, default_rtol_tensor, static_cast<T>(0.0));
+
     DenseTensor atol_compare_result;
-    atol_compare_result.Resize(atol.dims());
+    atol_compare_result.Resize(default_rtol_tensor.dims());
     phi::EqualKernel<T, Context>(
         dev_ctx, atol, zero_tensor, &atol_compare_result);
 
     DenseTensor selected_rtol_tensor;
-    selected_rtol_tensor.Resize(atol.dims());
+    selected_rtol_tensor.Resize(default_rtol_tensor.dims());
     phi::WhereKernel<T, Context>(dev_ctx,
                                  atol_compare_result,
                                  default_rtol_tensor,
diff --git a/paddle/phi/kernels/gpu/matrix_rank_tol_kernel.cu b/paddle/phi/kernels/gpu/matrix_rank_tol_kernel.cu
@@ -544,23 +544,24 @@ void MatrixRankAtolRtolKernel(const Context& dev_ctx,
     // when `rtol` is specified to be None in py api
     // use rtol=eps*max(m, n) only if `atol` is passed with value 0.0, else use
     // rtol=0.0
-    DenseTensor zero_tensor;
-    zero_tensor = phi::FullLike<T, Context>(dev_ctx, atol, static_cast<T>(0.0));
-
     T rtol_T = std::numeric_limits<T>::epsilon() * std::max(rows, cols);
     DenseTensor default_rtol_tensor;
     default_rtol_tensor =
         phi::Full<T, Context>(dev_ctx, {1}, static_cast<T>(rtol_T));
     default_rtol_tensor =
         phi::Multiply<T>(dev_ctx, default_rtol_tensor, max_eigenvalue_tensor);
 
+    DenseTensor zero_tensor;
+    zero_tensor = phi::FullLike<T, Context>(
+        dev_ctx, default_rtol_tensor, static_cast<T>(0.0));
+
     DenseTensor atol_compare_result;
-    atol_compare_result.Resize(atol.dims());
+    atol_compare_result.Resize(default_rtol_tensor.dims());
     phi::EqualKernel<T, Context>(
         dev_ctx, atol, zero_tensor, &atol_compare_result);
 
     DenseTensor selected_rtol_tensor;
-    selected_rtol_tensor.Resize(atol.dims());
+    selected_rtol_tensor.Resize(default_rtol_tensor.dims());
     phi::WhereKernel<T, Context>(dev_ctx,
                                  atol_compare_result,
                                  default_rtol_tensor,
diff --git a/test/legacy_test/test_matrix_rank_atol_rtol_op.py b/test/legacy_test/test_matrix_rank_atol_rtol_op.py
@@ -41,7 +41,7 @@ def np_matrix_rank_atol_rtol(x, atol=None, rtol=None, hermitian=False):
     sv_max = sv.max(axis=-1)
     if use_default_tol:
         rtol_T = np.finfo(x.dtype).eps * max(x.shape[-2], x.shape[-1]) * sv_max
-        rtol_default = np.full(rtol.shape, rtol_T, x.dtype)
+        rtol_default = np.full(x.shape[:-2], rtol_T, x.dtype)
         rtol = np.where(atol == 0.0, rtol_default, rtol * sv_max)
     else:
         rtol = rtol * sv_max
@@ -165,6 +165,17 @@ def test_dygraph(self):
             )
             np.testing.assert_allclose(rank_np, rank_pd, rtol=1e-05)
 
+            # atol: float, rtol: None
+            x_np = np.ones([3, 4, 5, 5])
+            x_pd = paddle.to_tensor(x_np)
+            rank_np = np_matrix_rank_atol_rtol(
+                x_np, atol=0.5, rtol=None, hermitian=True
+            )
+            rank_pd = paddle.linalg.matrix_rank(
+                x_pd, hermitian=True, atol=0.5, rtol=None
+            )
+            np.testing.assert_allclose(rank_np, rank_pd, rtol=1e-05)
+
             # atol: None, rtol: float
             x_np = np.eye(10, dtype=np.float32)
             x_pd = paddle.to_tensor(x_np)
@@ -307,6 +318,26 @@ def test_static(self):
                 )
                 np.testing.assert_allclose(fetches[0], rank_np, rtol=1e-05)
 
+        for place in places:
+            # atol: float, rtol: None
+            with static.program_guard(static.Program(), static.Program()):
+                x_np = np.ones([3, 4, 5, 5])
+                x_pd = paddle.static.data(
+                    name="X", shape=x_np.shape, dtype='float64'
+                )
+                rank_np = np_matrix_rank_atol_rtol(
+                    x_np, atol=0.5, rtol=None, hermitian=True
+                )
+                rank_pd = paddle.linalg.matrix_rank(
+                    x_pd, hermitian=True, atol=0.5, rtol=None
+                )
+                exe = base.Executor(place)
+                fetches = exe.run(
+                    feed={"X": x_np},
+                    fetch_list=[rank_pd],
+                )
+                np.testing.assert_allclose(fetches[0], rank_np, rtol=1e-05)
+
         for place in places:
             # atol: float, rtol: float
             with static.program_guard(static.Program(), static.Program()):
