[0-size Tensor No.76-79] Add 0-size Tensor support for i0 (#72826)

co63oc · web-flow · commit 0829b9896b08 · 2025-05-23T10:51:18.000+08:00
diff --git a/paddle/phi/kernels/cpu/i0_grad_kernel.cc b/paddle/phi/kernels/cpu/i0_grad_kernel.cc
@@ -26,6 +26,10 @@ void I0GradKernel(const Context& ctx,
                   const DenseTensor& x,
                   const DenseTensor& out_grad,
                   DenseTensor* x_grad) {
+  if (x_grad && x_grad->numel() == 0) {
+    ctx.template Alloc<T>(x_grad);
+    return;
+  }
   auto size = x.numel();
   auto* x_data = x.data<T>();
   auto* out_grad_data = out_grad.data<T>();
diff --git a/paddle/phi/kernels/cpu/i0_kernel.cc b/paddle/phi/kernels/cpu/i0_kernel.cc
@@ -23,6 +23,10 @@ namespace phi {
 
 template <typename T, typename Context>
 void I0Kernel(const Context& ctx, const DenseTensor& x, DenseTensor* out) {
+  if (out && out->numel() == 0) {
+    ctx.template Alloc<T>(out);
+    return;
+  }
   const int64_t size = x.numel();
   const T* x_data = x.data<T>();
   T* out_data = ctx.template Alloc<T>(out);
diff --git a/paddle/phi/kernels/cpu/i0e_grad_kernel.cc b/paddle/phi/kernels/cpu/i0e_grad_kernel.cc
@@ -27,6 +27,11 @@ void I0eGradKernel(const Context& ctx,
                    const DenseTensor& out,
                    const DenseTensor& out_grad,
                    DenseTensor* x_grad) {
+  if (x_grad && x_grad->numel() == 0) {
+    ctx.template Alloc<T>(x_grad);
+    return;
+  }
+
   auto size = x.numel();
   auto* x_data = x.data<T>();
   auto* out_data = out.data<T>();
diff --git a/paddle/phi/kernels/cpu/i0e_kernel.cc b/paddle/phi/kernels/cpu/i0e_kernel.cc
@@ -23,6 +23,10 @@ namespace phi {
 
 template <typename T, typename Context>
 void I0eKernel(const Context& ctx, const DenseTensor& x, DenseTensor* out) {
+  if (out && out->numel() == 0) {
+    ctx.template Alloc<T>(out);
+    return;
+  }
   int64_t size = x.numel();
   const T* x_data = x.data<T>();
   T* out_data = ctx.template Alloc<T>(out);
diff --git a/paddle/phi/kernels/cpu/i1_grad_kernel.cc b/paddle/phi/kernels/cpu/i1_grad_kernel.cc
@@ -27,6 +27,10 @@ void I1GradKernel(const Context& ctx,
                   const DenseTensor& out,
                   const DenseTensor& out_grad,
                   DenseTensor* x_grad) {
+  if (x_grad && x_grad->numel() == 0) {
+    ctx.template Alloc<T>(x_grad);
+    return;
+  }
   const int64_t size = x.numel();
   const T* x_data = x.data<T>();
   const T* out_data = out.data<T>();
diff --git a/paddle/phi/kernels/cpu/i1_kernel.cc b/paddle/phi/kernels/cpu/i1_kernel.cc
@@ -23,6 +23,10 @@ namespace phi {
 
 template <typename T, typename Context>
 void I1Kernel(const Context& ctx, const DenseTensor& x, DenseTensor* out) {
+  if (out && out->numel() == 0) {
+    ctx.template Alloc<T>(out);
+    return;
+  }
   const int64_t size = x.numel();
   const T* x_data = x.data<T>();
   T* out_data = ctx.template Alloc<T>(out);
diff --git a/paddle/phi/kernels/cpu/i1e_grad_kernel.cc b/paddle/phi/kernels/cpu/i1e_grad_kernel.cc
@@ -27,6 +27,10 @@ void I1eGradKernel(const Context& ctx,
                    const DenseTensor& out,
                    const DenseTensor& out_grad,
                    DenseTensor* x_grad) {
+  if (x_grad && x_grad->numel() == 0) {
+    ctx.template Alloc<T>(x_grad);
+    return;
+  }
   const int64_t size = x.numel();
   const T* x_data = x.data<T>();
   const T* out_data = out.data<T>();
diff --git a/paddle/phi/kernels/cpu/i1e_kernel.cc b/paddle/phi/kernels/cpu/i1e_kernel.cc
@@ -23,6 +23,10 @@ namespace phi {
 
 template <typename T, typename Context>
 void I1eKernel(const Context& ctx, const DenseTensor& x, DenseTensor* out) {
+  if (out && out->numel() == 0) {
+    ctx.template Alloc<T>(out);
+    return;
+  }
   const int64_t size = x.numel();
   const T* x_data = x.data<T>();
   T* out_data = ctx.template Alloc<T>(out);
diff --git a/test/legacy_test/test_i0_op.py b/test/legacy_test/test_i0_op.py
@@ -148,5 +148,26 @@ def test_check_grad(self):
         )
 
 
+class TestI0Op_ZeroSize(OpTest):
+    def setUp(self) -> None:
+        self.__class__.op_type = "i0"
+        self.op_type = "i0"
+        self.python_api = paddle.i0
+        self.init_config()
+        x = np.random.randn(3, 4, 0)
+        self.inputs = {'x': x.astype(self.dtype)}
+        self.attrs = {}
+        self.outputs = {'out': special.i0(x)}
+
+    def init_config(self):
+        self.dtype = np.float32
+
+    def test_check_output(self):
+        self.check_output()
+
+    def test_check_grad(self):
+        self.check_grad(['x'], 'out')
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/legacy_test/test_i0e_op.py b/test/legacy_test/test_i0e_op.py
@@ -152,5 +152,26 @@ def test_check_grad(self):
         )
 
 
+class TestI0eOp_ZeroSize(OpTest):
+    def setUp(self) -> None:
+        self.__class__.op_type = "i0e"
+        self.op_type = "i0e"
+        self.python_api = paddle.i0e
+        self.init_config()
+        x = np.random.randn(3, 4, 0)
+        self.inputs = {'x': x.astype(self.dtype)}
+        self.attrs = {}
+        self.outputs = {'out': special.i0e(x)}
+
+    def init_config(self):
+        self.dtype = np.float32
+
+    def test_check_output(self):
+        self.check_output()
+
+    def test_check_grad(self):
+        self.check_grad(['x'], 'out')
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/legacy_test/test_i1_op.py b/test/legacy_test/test_i1_op.py
@@ -156,5 +156,26 @@ def init_config(self):
         self.target = reference_i1(self.inputs['x'])
 
 
+class TestI1Op_ZeroSize(OpTest):
+    def setUp(self) -> None:
+        self.__class__.op_type = "i1"
+        self.op_type = "i1"
+        self.python_api = paddle.i1
+        self.init_config()
+        x = np.random.randn(3, 4, 0)
+        self.inputs = {'x': x.astype(self.dtype)}
+        self.attrs = {}
+        self.outputs = {'out': special.i1(x)}
+
+    def init_config(self):
+        self.dtype = np.float32
+
+    def test_check_output(self):
+        self.check_output()
+
+    def test_check_grad(self):
+        self.check_grad(['x'], 'out')
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/legacy_test/test_i1e_op.py b/test/legacy_test/test_i1e_op.py
@@ -156,5 +156,26 @@ def init_config(self):
         self.target = reference_i1e(self.inputs['x'])
 
 
+class TestI1eOp_ZeroSize(OpTest):
+    def setUp(self) -> None:
+        self.__class__.op_type = "i1e"
+        self.op_type = "i1e"
+        self.python_api = paddle.i1e
+        self.init_config()
+        x = np.random.randn(3, 4, 0)
+        self.inputs = {'x': x.astype(self.dtype)}
+        self.attrs = {}
+        self.outputs = {'out': special.i1e(x)}
+
+    def init_config(self):
+        self.dtype = np.float32
+
+    def test_check_output(self):
+        self.check_output()
+
+    def test_check_grad(self):
+        self.check_grad(['x'], 'out')
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/white_list/op_accuracy_white_list.py b/test/white_list/op_accuracy_white_list.py
@@ -79,6 +79,10 @@
     'fake_channel_wise_quantize_dequantize_abs_max',
     'fake_quantize_dequantize_abs_max',
     'fake_quantize_dequantize_moving_average_abs_max',
+    'i0',
+    'i1',
+    'i0e',
+    'i1e',
 ]
 
 NO_FP16_CHECK_GRAD_OP_LIST = [
