Use gather to replace gather nd when the bool index is one-dimensional (#72625)

zhanghonggeng · web-flow · commit f51e3ff85d10 · 2025-05-12T11:03:01.000+08:00
diff --git a/paddle/fluid/pybind/slice_utils.h b/paddle/fluid/pybind/slice_utils.h
@@ -519,6 +519,12 @@ static paddle::Tensor getValueForBoolTensor(const paddle::Tensor& tensor,
   if (bool_index.shape().size() == tensor_shape.size()) {
     return masked_select_ad_func(tensor, bool_index);
   }
+
+  if (bool_index.shape().size() == 1) {
+    auto bool_2_idx = nonzero_ad_func(bool_index);
+    return gather_ad_func(tensor, bool_2_idx);
+  }
+
   auto bool_2_idx = nonzero_ad_func(bool_index);
   return gather_nd_ad_func(tensor, bool_2_idx);
 }

Original file line number	Diff line number	Diff line change
`@@ -519,6 +519,12 @@ static paddle::Tensor getValueForBoolTensor(const paddle::Tensor& tensor,`
`519`	`519`	`if (bool_index.shape().size() == tensor_shape.size()) {`
`520`	`520`	`return masked_select_ad_func(tensor, bool_index);`
`521`	`521`	`}`
	`522`	`+`
	`523`	`+ if (bool_index.shape().size() == 1) {`
	`524`	`+ auto bool_2_idx = nonzero_ad_func(bool_index);`
	`525`	`+ return gather_ad_func(tensor, bool_2_idx);`
	`526`	`+ }`
	`527`	`+`
`522`	`528`	`auto bool_2_idx = nonzero_ad_func(bool_index);`
`523`	`529`	`return gather_nd_ad_func(tensor, bool_2_idx);`
`524`	`530`	`}`