【HEU】[Paddle Tensor 第二期 API支持 0-size Tensor] paddle.cross 支持 0-size tensor (#70103)

* fix cross

* fix

* fix test

---------

Co-authored-by: fangfangssj <[email protected]>

Author: fangfangssj

paddle/phi/kernels/cpu/cross_kernel.cc

@@ -70,6 +70,11 @@ void CrossKernel(const Context& dev_ctx,
                           "But received: Input(X/Y).dims() == [%s].",
                           input_x_dims));
   }
+  if (input_x.numel() == 0 || input_y.numel() == 0) {
+    output->Resize(input_x.dims());
+    dev_ctx.template Alloc<T>(output);
+    return;
+  }
   auto outer_loops = 1;
   for (auto i = 0; i < dim; i++) {
     outer_loops *= static_cast<int>(input_x_dims[i]);

paddle/phi/kernels/gpu/cross_kernel.cu

@@ -101,6 +101,11 @@ void CrossKernel(const Context& dev_ctx,
                           input_x_dims));
   }
 
+  if (input_x.numel() == 0 || input_y.numel() == 0) {
+    output->Resize(input_x.dims());
+    dev_ctx.template Alloc<T>(output);
+    return;
+  }
   std::vector<int> cal_dims;
   std::vector<int> left_strides;
   std::vector<int> full_strides;

test/legacy_test/test_cross_op.py

@@ -170,6 +170,8 @@ def input_data(self):
         self.data_y = np.array(
             [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]
         ).astype('float32')
+        self.data_x_zero = np.array([]).reshape(0, 3).astype('float32')
+        self.data_y_zero = np.array([]).reshape(0, 3).astype('float32')
 
     def test_cross_api(self):
         self.input_data()
@@ -212,6 +214,26 @@ def test_cross_api(self):
         )
         np.testing.assert_allclose(expect_out, np.array(res), rtol=1e-05)
 
+        main = paddle.static.Program()
+        startup = paddle.static.Program()
+        # case 3:
+        with paddle.static.program_guard(main, startup):
+            x = paddle.static.data(name='x', shape=[0, 3], dtype="float32")
+            y = paddle.static.data(name='y', shape=[0, 3], dtype="float32")
+            z = paddle.cross(x, y, axis=1)
+            exe = base.Executor(base.CPUPlace())
+            (res,) = exe.run(
+                main,
+                feed={'x': self.data_x_zero, 'y': self.data_y_zero},
+                fetch_list=[z],
+                return_numpy=False,
+            )
+        expect_out = np.empty((0, 3))
+        np.testing.assert_allclose(expect_out, np.array(res), rtol=1e-05)
+
+        main = paddle.static.Program()
+        startup = paddle.static.Program()
+
     def test_cross_api1(self):
         with paddle.pir_utils.OldIrGuard():
             self.input_data()
@@ -227,6 +249,17 @@ def test_cross_api1(self):
                 y_1 = paddle.cross(x, y, name='result')
                 self.assertEqual(('result' in y_1.name), True)
 
+            main = paddle.static.Program()
+            startup = paddle.static.Program()
+
+            # case 2:
+            with paddle.static.program_guard(main, startup):
+                x = paddle.static.data(name="x", shape=[0, 3], dtype="float32")
+                y = paddle.static.data(name='y', shape=[0, 3], dtype='float32')
+
+                y_1 = paddle.cross(x, y, axis=1, name='result')
+                self.assertEqual(('result' in y_1.name), True)
+
     def test_dygraph_api(self):
         self.input_data()
         # case 1:
@@ -250,6 +283,15 @@ def test_dygraph_api(self):
         )
         np.testing.assert_allclose(expect_out, np_z, rtol=1e-05)
 
+        # case 3:
+        with base.dygraph.guard():
+            x = paddle.to_tensor(self.data_x_zero)
+            y = paddle.to_tensor(self.data_y_zero)
+            z = paddle.cross(x, y, axis=1)
+            np_z = z.numpy()
+        expect_out = np.empty((0, 3))
+        np.testing.assert_allclose(expect_out, np_z, rtol=1e-05)
+
 
 if __name__ == '__main__':
     paddle.enable_static()