【Paddle Tensor 第二期 API支持 0-size Tensor】Fix frobenius_norm to support 0-size tensor

paddle/phi/kernels/gpu/frobenius_norm_kernel.cu

@@ -16,6 +16,7 @@
 
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/activation_kernel.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/activation_functor.h"
 #include "paddle/phi/kernels/gpu/reduce.h"
 
@@ -28,6 +29,28 @@ void FrobeniusNormKernel(const Context& dev_ctx,
                          bool keep_dim,
                          bool reduce_all,
                          DenseTensor* out) {
+  if (x.numel() == 0) {
+    auto dim_x = x.dims();
+    std::set<int> axis_set;
+    for (auto ax : dims.GetData()) {
+      if (ax < 0) ax += dim_x.size();
+      axis_set.insert(ax);
+    }
+
+    std::vector<int64_t> out_dims_vec;
+    for (int i = 0; i < dim_x.size(); ++i) {
+      if (axis_set.count(i) == 0) {
+        out_dims_vec.push_back(dim_x[i]);
+      } else if (keep_dim) {
+        out_dims_vec.push_back(1);
+      }
+    }
+
+    // 正确调用 Full 初始化
+    phi::Full<T, Context>(
+        dev_ctx, phi::IntArray(out_dims_vec), static_cast<T>(0), out);
+    return;
+  }
   reduce_all = recompute_reduce_all(x, dims.GetData(), reduce_all);
   auto out_dtype = x.dtype();
   phi::Reduce<T, kps::AddFunctor, kps::SquareFunctor>(

paddle/phi/kernels/impl/frobenius_norm_kernel_impl.h

@@ -27,6 +27,30 @@ void FrobeniusNormKernel(const Context& ctx,
                          bool keep_dim,
                          bool reduce_all,
                          DenseTensor* out) {
+  auto xdim = x.dims();
+
+  if (x.numel() == 0) {
+    std::set<int> axis_set;
+    for (int ax : axis.GetData()) {
+      if (ax < 0) {
+        ax += xdim.size();
+      }
+      axis_set.insert(ax);
+    }
+
+    std::vector<int64_t> out_dims_vec;
+    for (int i = 0; i < xdim.size(); ++i) {
+      if (axis_set.find(i) == axis_set.end()) {
+        out_dims_vec.push_back(xdim[i]);
+      } else if (keep_dim) {
+        out_dims_vec.push_back(1);
+      }
+    }
+    out->Resize(phi::make_ddim(out_dims_vec));
+    ctx.template Alloc<T>(out);
+    phi::funcs::SetConstant<Context, T>()(ctx, out, 0);
+    return;
+  }
   reduce_all = recompute_reduce_all(x, axis.GetData(), reduce_all);
   Reduce<Context, T, funcs::FrobeniusNormFunctor>(
       ctx, x, reduce_all, axis.GetData(), keep_dim, x.dtype(), out);

test/legacy_test/test_norm_all.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 import unittest
 
 import numpy as np
@@ -182,6 +183,49 @@ 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]
+        self.axis = (1, 2)
+        self.keepdim = False
+
+    def init_dtype(self):
+        self.dtype = "float32"
+
+    def test_check_output(self):
+        places = (
+            [paddle.CPUPlace(), paddle.CUDAPlace(0)]
+            if core.is_compiled_with_cuda()
+            else [paddle.CPUPlace()]
+        )
+        for place in places:
+            self.check_output_with_place(place)
+
+    def test_check_grad(self):
+        pass
+
+
+class TestFrobeniusNormOpZeroSize2(TestFrobeniusNormOpZeroSize):
+    def init_test_case(self):
+        self.shape = [3, 0, 3]
+        self.axis = (1, 2)
+        self.keepdim = False
+
+
+class TestFrobeniusNormOpZeroSize3(TestFrobeniusNormOpZeroSize):
+    def init_test_case(self):
+        self.shape = [0, 20, 3]
+        self.axis = (0, 2)
+        self.keepdim = False
+
+
+class TestFrobeniusNormOpZeroSize4(TestFrobeniusNormOpZeroSize):
+    def init_test_case(self):
+        self.shape = [0, 20, 3]
+        self.axis = (0, -1)
+        self.keepdim = False
+
+
 class TestPnormOp(OpTest):
     def setUp(self):
         self.op_type = "p_norm"
@@ -690,6 +734,7 @@ def check_linalg_vector_dygraph(
 
 class API_NormTest(unittest.TestCase):
     def test_basic(self):
+        paddle.enable_static()
         keep_dims = {False, True}
         for keep in keep_dims:
             check_fro_static(
@@ -1478,6 +1523,96 @@ def err_dtype(p, shape_x, xdtype, out=None):
             )
 
 
+class TestMatrixNormZeroSizeTensorTensor(unittest.TestCase):
+    def _get_places(self):
+        places = []
+        if (
+            os.environ.get('FLAGS_CI_both_cpu_and_gpu', 'False').lower()
+            in ['1', 'true', 'on']
+            or not paddle.is_compiled_with_cuda()
+        ):
+            places.append(base.CPUPlace())
+        if paddle.is_compiled_with_cuda():
+            places.append(base.CUDAPlace(0))
+        return places
+
+    def _test_matrix_norm_static(self, place):
+        with paddle.static.program_guard(
+            paddle.static.Program(), paddle.static.Program()
+        ):
+            x1 = paddle.static.data(name='x1', shape=[0, 0], dtype='float32')
+            x2 = paddle.static.data(name='x2', shape=[2, 2, 0], dtype='float32')
+            x3 = paddle.static.data(
+                name='x3', shape=[3, 0, 2, 0], dtype='float32'
+            )
+
+            y1_fro = paddle.linalg.matrix_norm(x1, p='fro')
+            y2_fro = paddle.linalg.matrix_norm(x2, p='fro')
+            y3_fro = paddle.linalg.matrix_norm(x3, p='fro', axis=[1, -1])
+
+            y2_p1 = paddle.linalg.matrix_norm(x2, p=1, axis=[0, 1])
+            y2_p2 = paddle.linalg.matrix_norm(x2, p=2, axis=[0, 1])
+            y2_pinf = paddle.linalg.matrix_norm(x2, p=np.inf, axis=[0, 1])
+            y2_pninf = paddle.linalg.matrix_norm(x2, p=-np.inf, axis=[0, 1])
+
+            exe = paddle.static.Executor(place)
+            fetch_list = [
+                y1_fro,
+                y2_fro,
+                y3_fro,
+                y2_p1,
+                y2_p2,
+                y2_pinf,
+                y2_pninf,
+            ]
+            res = exe.run(
+                feed={
+                    'x1': np.zeros((0, 0), dtype='float32'),
+                    'x2': np.ones((2, 2, 0), dtype='float32'),
+                    'x3': np.ones((3, 0, 2, 0), dtype='float32'),
+                },
+                fetch_list=fetch_list,
+            )
+
+            self.assertEqual(res[0].shape, ())  # y1_fro
+            self.assertEqual(res[1].shape, (2,))  # y2_fro
+            self.assertEqual(res[2].shape, (3, 2))  # y3_fro
+            for r in res[3:]:
+                self.assertEqual(r.shape, (0,))
+
+    def _test_matrix_norm_dynamic(self):
+        paddle.disable_static()
+        for place in self._get_places():
+            paddle.set_device(
+                'gpu' if isinstance(place, paddle.CUDAPlace) else 'cpu'
+            )
+
+            x1 = paddle.full((0, 0), 1.0, dtype='float32')
+            x2 = paddle.full((2, 2, 0), 1.0, dtype='float32')
+            x3 = paddle.full((3, 0, 2, 0), 1.0, dtype='float32')
+
+            y1_fro = paddle.linalg.matrix_norm(x1, p='fro')
+            y2_fro = paddle.linalg.matrix_norm(x2, p='fro')
+            y3_fro = paddle.linalg.matrix_norm(x3, p='fro', axis=[1, -1])
+
+            y2_p1 = paddle.linalg.matrix_norm(x2, p=1, axis=[0, 1])
+            y2_p2 = paddle.linalg.matrix_norm(x2, p=2, axis=[0, 1])
+            y2_pinf = paddle.linalg.matrix_norm(x2, p=np.inf, axis=[0, 1])
+            y2_pninf = paddle.linalg.matrix_norm(x2, p=-np.inf, axis=[0, 1])
+
+            self.assertEqual(y1_fro.shape, [])
+            self.assertEqual(y2_fro.shape, (2,))
+            self.assertEqual(y3_fro.shape, (3, 2))
+            for y in [y2_p1, y2_p2, y2_pinf, y2_pninf]:
+                self.assertEqual(y.shape, (0,))
+
+    def test_matrix_norm_zero_size(self):
+        for place in self._get_places():
+            self._test_matrix_norm_static(place)
+        self._test_matrix_norm_dynamic()
+        paddle.enable_static()
+
+
 if __name__ == '__main__':
     paddle.enable_static()
     unittest.main()