[0-size Tensor Job2 No.58] Add 0-size Tensor support for gather_tree (#73387)

co63oc · web-flow · commit 864b52061557 · 2025-06-19T14:08:37.000+08:00
* Fix

* Fix
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc
@@ -1058,6 +1058,9 @@ bool GatherTreeOpInferSymbolicShape(
                         "shape of Input(Ids)."));
   size_t rank = ids_shape.size();
   for (size_t i = 0; i < rank; ++i) {
+    if (ids_shape[i] == 0) {
+      continue;
+    }
     infer_context->AddEqualCstr(ids_shape[i], parents_shape[i]);
   }
 
diff --git a/paddle/phi/infermeta/binary.cc b/paddle/phi/infermeta/binary.cc
@@ -2138,11 +2138,13 @@ void GatherTreeMeta(const MetaTensor& ids,
                     MetaTensor* out) {
   auto ids_dims = ids.dims();
   auto parents_dims = parents.dims();
-  PADDLE_ENFORCE_EQ(ids_dims == parents_dims,
-                    true,
-                    common::errors::InvalidArgument(
-                        "The shape of Input(Parents) must be same with the "
-                        "shape of Input(Ids)."));
+  if (common::product(ids_dims) != 0) {
+    PADDLE_ENFORCE_EQ(ids_dims == parents_dims,
+                      true,
+                      common::errors::InvalidArgument(
+                          "The shape of Input(Parents) must be same with the "
+                          "shape of Input(Ids)."));
+  }
   out->set_dims(ids_dims);
   out->set_dtype(ids.dtype());
 }
diff --git a/paddle/phi/kernels/cpu/gather_tree_kernel.cc b/paddle/phi/kernels/cpu/gather_tree_kernel.cc
@@ -28,6 +28,9 @@ void GatherTreeKernel(const Context &dev_ctx,
   const auto *parents_data = parents.data<T>();
 
   T *out_data = dev_ctx.template Alloc<T>(out);
+  if (out && out->numel() == 0) {
+    return;
+  }
 
   auto &ids_dims = ids.dims();
   int64_t max_length = ids_dims[0];
diff --git a/paddle/phi/kernels/gpu/gather_tree_kernel.cu b/paddle/phi/kernels/gpu/gather_tree_kernel.cu
@@ -63,6 +63,9 @@ void GatherTreeKernel(const Context &dev_ctx,
   const auto *ids_data = ids.data<T>();
   const auto *parents_data = parents.data<T>();
   T *out_data = dev_ctx.template Alloc<T>(out);
+  if (out && out->numel() == 0) {
+    return;
+  }
 
   PADDLE_ENFORCE_NOT_NULL(ids_data,
                           common::errors::InvalidArgument(
diff --git a/test/legacy_test/test_gather_tree_op.py b/test/legacy_test/test_gather_tree_op.py
@@ -150,6 +150,33 @@ def test_parents_ndim():
         paddle.disable_static()
 
 
+class TestGatherTreeOp_ZeroSize(OpTest):
+    def init_shape(self):
+        self.ids_shape = (0, 2, 2)
+        self.parents_shape = (0, 2, 2)
+
+    def setUp(self):
+        self.op_type = "gather_tree"
+        self.python_api = paddle.nn.functional.gather_tree
+        self.init_shape()
+        ids_shape = self.ids_shape
+        parents_shape = self.parents_shape
+        max_length, batch_size, beam_size = ids_shape
+        ids = np.random.randint(0, high=10, size=ids_shape)
+        parents = np.random.randint(0, high=beam_size, size=parents_shape)
+        self.inputs = {"Ids": ids, "Parents": parents}
+        self.outputs = {'Out': ids}
+
+    def test_check_output(self):
+        self.check_output(check_pir=True)
+
+
+class TestGatherTreeOp_ZeroSize2(TestGatherTreeOp_ZeroSize):
+    def init_shape(self):
+        self.ids_shape = (0, 2, 2)
+        self.parents_shape = (1, 2, 2)
+
+
 if __name__ == "__main__":
     paddle.enable_static()
     unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -1058,6 +1058,9 @@ bool GatherTreeOpInferSymbolicShape(`
`1058`	`1058`	`"shape of Input(Ids)."));`
`1059`	`1059`	`size_t rank = ids_shape.size();`
`1060`	`1060`	`for (size_t i = 0; i < rank; ++i) {`
	`1061`	`+ if (ids_shape[i] == 0) {`
	`1062`	`+ continue;`
	`1063`	`+ }`
`1061`	`1064`	`infer_context->AddEqualCstr(ids_shape[i], parents_shape[i]);`
`1062`	`1065`	`}`
`1063`	`1066`