[0-size Tensor Job2 No.17-20、31-34] Add 0-size Tensor support for segment_max (#73341)

co63oc · web-flow · commit e636f35c603e · 2025-06-17T12:10:51.000+08:00
* Fix

* Fix
diff --git a/paddle/phi/kernels/impl/segment_pool_grad_kernel_impl.h b/paddle/phi/kernels/impl/segment_pool_grad_kernel_impl.h
@@ -34,6 +34,10 @@ void SegmentPoolGradKernel(const Context& dev_ctx,
   dev_ctx.template Alloc<T>(x_grad);
   phi::funcs::SetConstant<Context, T> set_zero;
   set_zero(dev_ctx, x_grad, static_cast<T>(0));
+  // return after allocation if x_grad is empty.
+  if (x_grad && x_grad->numel() == 0) {
+    return;
+  }
 
   auto index_type = segment_ids.type();
   if (index_type == DataType::INT32) {
diff --git a/paddle/phi/kernels/impl/segment_pool_kernel_impl.h b/paddle/phi/kernels/impl/segment_pool_kernel_impl.h
@@ -42,10 +42,6 @@ void SegmentKernelLaunchHelper(const Context& dev_ctx,
                         "dimension size is 1. Segment_ids's shape is: [%s].",
                         segment_ids.dims()));
 
-  if (x.numel() == 0 || segment_ids.numel() == 0) {
-    return;
-  }
-
   bool cpu_place = dev_ctx.GetPlace().GetType() == phi::AllocationType::CPU;
   if (cpu_place) {
     auto dims = x.dims();
@@ -112,6 +108,10 @@ void SegmentKernelLaunchHelper(const Context& dev_ctx,
     }
   }
 #endif
+  // return after allocation, if x or segment_ids is empty tensor.
+  if (x.numel() == 0 || segment_ids.numel() == 0) {
+    return;
+  }
 
   phi::funcs::SegmentPoolFunctor<Context, T, IndexT> pool;
 
diff --git a/test/legacy_test/test_segment_ops.py b/test/legacy_test/test_segment_ops.py
@@ -357,6 +357,79 @@ def test_check_grad(self):
         self.check_grad_with_place(self.place, ["X"], "Out", check_pir=True)
 
 
+# default SUM
+class TestSegmentOps_ZeroSize(OpTest):
+    def set_data(self):
+        x = np.random.uniform(-1, 1, self.shape).astype(self.dtype)
+        segment_ids = self.set_segment(self.shape[0])
+        return x, segment_ids
+
+    def set_segment(self, len):
+        segment = np.random.randint(0, len, size=[len])
+        segment = np.sort(segment)
+        return segment.astype('int64')
+
+    def compute(self, x, segment_ids):
+        return compute_segment_sum(x, segment_ids)
+
+    def prepare(self):
+        self.op_type = "segment_pool"
+        self.python_api = segment_pool_split
+        self.python_out_sig = ["Out"]
+        self.dtype = np.float64
+        self.shape = [30, 0]
+        self.attrs = {"pooltype": "SUM"}
+
+    def setUp(self):
+        self.prepare()
+        x, segment_ids = self.set_data()
+        result = self.compute(x, segment_ids)
+        self.inputs = {
+            'X': x,
+            'SegmentIds': segment_ids.astype(np.int64),
+        }
+        self.outputs = {'Out': result.astype(self.dtype)}
+
+    def test_check_output(self):
+        self.check_output(check_pir=True, check_symbol_infer=False)
+
+    def test_check_grad(self):
+        self.check_grad(["X"], "Out", check_pir=True)
+
+
+class TestSegmentOps_Min_ZeroSize(TestSegmentOps_ZeroSize):
+    def compute(self, x, segment_ids):
+        result, self.gradient = compute_segment_min_max(
+            x, segment_ids, pooltype="MIN"
+        )
+        return result
+
+    def prepare(self):
+        super().prepare()
+        self.attrs = {'pooltype': "MIN"}
+
+
+class TestSegmentOps_Max_ZeroSize(TestSegmentOps_ZeroSize):
+    def compute(self, x, segment_ids):
+        result, self.gradient = compute_segment_min_max(
+            x, segment_ids, pooltype="MAX"
+        )
+        return result
+
+    def prepare(self):
+        super().prepare()
+        self.attrs = {'pooltype': "MAX"}
+
+
+class TestSegmentOps_Mean_ZeroSize(TestSegmentOps_ZeroSize):
+    def compute(self, x, segment_ids):
+        return compute_segment_mean(x, segment_ids)
+
+    def prepare(self):
+        super().prepare()
+        self.attrs = {'pooltype': "MEAN"}
+
+
 class API_SegmentOpsTest(unittest.TestCase):
 
     def test_static(self):
