frobenius_norm support complex input (#72577)

aquagull · web-flow · commit fa0f36fd715f · 2025-06-09T15:20:53.000+08:00
* fix p

* fix _backshift_permutation

* update

* update

* fix

* revert kernel change, code-style

* fix

* fix

* fix

* fix

* fix

* add test
diff --git a/paddle/phi/kernels/cpu/frobenius_norm_grad_kernel.cc b/paddle/phi/kernels/cpu/frobenius_norm_grad_kernel.cc
@@ -22,4 +22,6 @@ PD_REGISTER_KERNEL(frobenius_norm_grad,
                    ALL_LAYOUT,
                    phi::FrobeniusNormGradKernel,
                    float,
-                   double) {}
+                   double,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
diff --git a/paddle/phi/kernels/cpu/frobenius_norm_kernel.cc b/paddle/phi/kernels/cpu/frobenius_norm_kernel.cc
@@ -17,5 +17,11 @@
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/impl/frobenius_norm_kernel_impl.h"
 
-PD_REGISTER_KERNEL(
-    frobenius_norm, CPU, ALL_LAYOUT, phi::FrobeniusNormKernel, float, double) {}
+PD_REGISTER_KERNEL(frobenius_norm,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::FrobeniusNormKernel,
+                   float,
+                   double,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
diff --git a/paddle/phi/kernels/gpu/frobenius_norm_grad_kernel.cu b/paddle/phi/kernels/gpu/frobenius_norm_grad_kernel.cu
@@ -22,4 +22,6 @@ PD_REGISTER_KERNEL(frobenius_norm_grad,
                    ALL_LAYOUT,
                    phi::FrobeniusNormGradKernel,
                    float,
-                   double) {}
+                   double,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
diff --git a/paddle/phi/kernels/gpu/frobenius_norm_kernel.cu b/paddle/phi/kernels/gpu/frobenius_norm_kernel.cu
@@ -43,5 +43,11 @@ void FrobeniusNormKernel(const Context& dev_ctx,
 
 }  // namespace phi
 
-PD_REGISTER_KERNEL(
-    frobenius_norm, GPU, ALL_LAYOUT, phi::FrobeniusNormKernel, float, double) {}
+PD_REGISTER_KERNEL(frobenius_norm,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::FrobeniusNormKernel,
+                   float,
+                   double,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
diff --git a/python/paddle/tensor/linalg.py b/python/paddle/tensor/linalg.py
@@ -773,8 +773,10 @@ def _backshift_permutation(dim0, dim1, dimn):
         Auxiliary function for matrix_norm
         Computes the permutation that moves the two given dimensions to the back
         """
-        ret = [i for i in range(dimn) if i != dim0 and i != dim1]
-        ret.extend((dim0, dim1))
+        pos_dim0 = dim0 % dimn
+        pos_dim1 = dim1 % dimn
+        ret = [i for i in range(dimn) if i != pos_dim0 and i != pos_dim1]
+        ret.extend((pos_dim0, pos_dim1))
         return ret
 
     def _inverse_permutation(perm):
@@ -792,7 +794,7 @@ def frobenius_norm(
         """
         The frobenius norm OP is to calculate the frobenius norm of certain two dimensions of Tensor `input`.
         Args:
-          input (Variable): Tensor, data type float32, float64.
+          input (Variable): Tensor, data type float32, float64, complex64, complex128.
           dim (list, optional): None for last two dimensions. Default None.
           keepdim (bool, optional): Whether keep the dimensions as the `input`, Default False.
           name (str, optional): The default value is None. Normally there is no need for
@@ -961,7 +963,8 @@ def p_matrix_norm(
                     )
                 return result
             else:  # 1,-1,inf,-inf
-                dim0, dim1 = axis
+                rank = len(x.shape)
+                dim0, dim1 = (d % rank for d in axis)
                 if abs_ord == np.float64("inf"):
                     dim0, dim1 = dim1, dim0
                 if not keepdim and (dim0 < dim1):
@@ -1056,7 +1059,8 @@ def p_matrix_norm(
             return reduce_out
 
         else:
-            dim0, dim1 = axis
+            rank = len(x.shape)
+            dim0, dim1 = (d % rank for d in axis)
             if abs_ord == np.float64("inf"):
                 dim0, dim1 = dim1, dim0
             if not keepdim and (dim0 < dim1):
diff --git a/test/legacy_test/test_norm_all.py b/test/legacy_test/test_norm_all.py
@@ -183,6 +183,32 @@ def test_check_grad(self):
         self.check_grad(['X'], 'Out', check_pir=True)
 
 
+class TestFrobeniusNormOp3(TestFrobeniusNormOp):
+    def init_test_case(self):
+        self.shape = [5, 5, 5]
+        self.axis = (0, 1)
+        self.keepdim = True
+
+    def init_dtype(self):
+        self.dtype = "complex64"
+
+    def test_check_grad(self):
+        self.check_grad(['X'], 'Out', check_pir=True)
+
+
+class TestFrobeniusNormOp4(TestFrobeniusNormOp):
+    def init_test_case(self):
+        self.shape = [5, 5, 5, 2]
+        self.axis = (0, 1)
+        self.keepdim = True
+
+    def init_dtype(self):
+        self.dtype = "complex128"
+
+    def test_check_grad(self):
+        self.check_grad(['X'], 'Out', check_pir=True)
+
+
 class TestFrobeniusNormOpZeroSize(TestFrobeniusNormOp):
     def init_test_case(self):
         self.shape = [0, 20, 3]
@@ -948,7 +974,6 @@ def test_basic(self):
                     keep_dim=keep,
                     check_dim=True,
                 )
-
                 check_linalg_vector_static(
                     self,
                     p=1,
@@ -1093,6 +1118,40 @@ def test_basic(self):
                     keep_dim=keep,
                     check_dim=True,
                 )
+                check_linalg_matrix_static(
+                    self,
+                    p=-np.inf,
+                    axis=[-2, -1],
+                    shape_x=[0, 1, 2, 1],
+                    dtype="float64",
+                    keep_dim=keep,
+                    check_dim=True,
+                )
+                check_linalg_matrix_static(
+                    self,
+                    p="fro",
+                    axis=[-2, -1],
+                    shape_x=[0, 1, 2, 1],
+                    dtype="float64",
+                    keep_dim=keep,
+                    check_dim=True,
+                )
+                check_linalg_matrix_static(
+                    self,
+                    p="fro",
+                    axis=[-2, -1],
+                    shape_x=[3, 2, 1],
+                    dtype="complex64",
+                    keep_dim=keep,
+                )
+                check_linalg_matrix_static(
+                    self,
+                    p="fro",
+                    axis=[-2, -1],
+                    shape_x=[3, 2, 1],
+                    dtype="complex128",
+                    keep_dim=keep,
+                )
 
     def test_dygraph(self):
         paddle.disable_static()
@@ -1135,7 +1194,6 @@ def test_dygraph(self):
                 keep_dim=keep,
                 check_dim=True,
             )
-
             check_linalg_norm_dygraph(
                 self,
                 p=2,
@@ -1466,7 +1524,40 @@ def test_dygraph(self):
                 keep_dim=keep,
                 check_dim=True,
             )
-
+            check_linalg_matrix_dygraph(
+                self,
+                p=-np.inf,
+                axis=[-2, -1],
+                shape_x=[0, 1, 2, 1],
+                dtype="float64",
+                keep_dim=keep,
+                check_dim=True,
+            )
+            check_linalg_matrix_dygraph(
+                self,
+                p="fro",
+                axis=[-2, -1],
+                shape_x=[0, 1, 2, 1],
+                dtype="float64",
+                keep_dim=keep,
+                check_dim=True,
+            )
+            check_linalg_matrix_dygraph(
+                self,
+                p="fro",
+                axis=[-2, -1],
+                shape_x=[3, 2, 1],
+                dtype="complex64",
+                keep_dim=keep,
+            )
+            check_linalg_matrix_dygraph(
+                self,
+                p="fro",
+                axis=[-2, -1],
+                shape_x=[3, 2, 1],
+                dtype="complex128",
+                keep_dim=keep,
+            )
         paddle.enable_static()
 
     def test_name(self):
