[0-size Tensor Job2 No.7、64、83] Add 0-size Tensor support for one_hot (#73386)

co63oc · web-flow · commit bc60082b216b · 2025-06-18T16:35:35.000+08:00
* Fix

* Fix

* Fix
diff --git a/paddle/phi/kernels/gpu/cumprod_grad_kernel.cu b/paddle/phi/kernels/gpu/cumprod_grad_kernel.cu
@@ -159,6 +159,11 @@ void CumprodGradKernel(const Context &dev_ctx,
   const auto *y = &out;
   const auto *dy = &dout;
 
+  if (dx && dx->numel() == 0) {
+    dev_ctx.template Alloc<T>(dx);
+    return;
+  }
+
   size_t outer_dim, mid_dim, inner_dim;
   GetCumprodDimInfo(x.dims(), dim, &outer_dim, &mid_dim, &inner_dim);
   if (x.dims().size() == 0) {
diff --git a/paddle/phi/kernels/gpu/one_hot_kernel.cu b/paddle/phi/kernels/gpu/one_hot_kernel.cu
@@ -57,6 +57,7 @@ void OneHotKernel(const Context& dev_ctx,
   auto* p_out_data = dev_ctx.template Alloc<float>(out);
   auto stream = dev_ctx.stream();
   funcs::set_constant(dev_ctx, out, static_cast<float>(0.0));
+  if (numel == 0) return;
 
   auto config = phi::backends::gpu::GetGpuLaunchConfig1D(dev_ctx, numel);
 
diff --git a/paddle/phi/kernels/xpu/one_hot_kernel.cc b/paddle/phi/kernels/xpu/one_hot_kernel.cc
@@ -33,6 +33,7 @@ void OneHotKernel(const Context& dev_ctx,
   auto* p_in_data = x.data<T>();
   auto numel = x.numel();
   auto* p_out_data = dev_ctx.template Alloc<float>(out);
+  if (numel == 0) return;
   int r = xpu::one_hot<T>(
       dev_ctx.x_context(), p_in_data, p_out_data, numel, depth_v, 1.0, 0.0);
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "one_hot");
diff --git a/test/legacy_test/test_cumprod_op.py b/test/legacy_test/test_cumprod_op.py
@@ -1091,5 +1091,40 @@ def test_check_grad(self):
                 )
 
 
+class TestCumprodAPI_ZeroSize(unittest.TestCase):
+    def init_dtype(self):
+        self.dtype = 'float64'
+        self.shape = [0, 3, 10, 10]
+
+    def setUp(self):
+        self.init_dtype()
+        self.x = (np.random.rand(0, 3, 10, 10) + 0.5).astype(self.dtype)
+        self.place = []
+        if (
+            os.environ.get('FLAGS_CI_both_cpu_and_gpu', 'False').lower()
+            in ['1', 'true', 'on']
+            or not core.is_compiled_with_cuda()
+        ):
+            self.place.append(paddle.CPUPlace())
+        if core.is_compiled_with_cuda():
+            self.place.append(paddle.CUDAPlace(0))
+
+    # test dynamic graph api.
+    def test_dygraph_api(self):
+        def run(place):
+            paddle.disable_static(place)
+            x = paddle.to_tensor(self.x)
+            x.stop_gradient = False
+            out = paddle.cumprod(x, 1)
+            out_ref = np.cumprod(self.x, 1)
+            np.testing.assert_allclose(out_ref, out.numpy(), rtol=1e-05)
+            paddle.sum(out).backward()
+            np.testing.assert_allclose(x.grad.shape, x.shape)
+            paddle.enable_static()
+
+        for place in self.place:
+            run(place)
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/legacy_test/test_one_hot_v2_op.py b/test/legacy_test/test_one_hot_v2_op.py
@@ -254,6 +254,36 @@ def test_bad_x():
             self.assertRaises(TypeError, test_bad_x)
 
 
+class TestOneHotOp_ZeroSize(OpTest):
+    def setUp(self):
+        self.op_type = 'one_hot_v2'
+        self.python_api = one_hot_wrapper
+        self.public_python_api = one_hot_wrapper
+        self.python_out_sig = ['Out']
+        depth = 10
+        depth_np = np.array(10).astype('int32')
+        x_shape = [0, 10, 7, 3]
+        x = [np.random.randint(0, depth - 1) for i in range(np.prod(x_shape))]
+        x = np.array(x).astype('int32').reshape(x_shape)
+
+        out = np.zeros(shape=(np.prod(x.shape), depth)).astype('float32')
+
+        r_x = np.reshape(x, np.prod(x.shape))
+        for i in range(np.prod(x.shape)):
+            out[i, r_x[i]] = 1.0
+
+        shape_np = list(x.shape)
+        shape_np.append(depth)
+        out = np.reshape(out, shape_np)
+
+        self.inputs = {'X': x, 'depth_tensor': depth_np}
+        self.attrs = {'dtype': int(core.VarDesc.VarType.FP32)}
+        self.outputs = {'Out': out}
+
+    def test_check_output(self):
+        self.check_output()
+
+
 if __name__ == '__main__':
     paddle.enable_static()
     unittest.main()
