Implement paddle.min/paddle.max for 0-size tensors along non-zero dimensions (#72565)

aquagull · web-flow · commit 14e9edc4bd74 · 2025-05-27T16:37:41.000+08:00
* update code

* update

* fix

* add onednn test

* trans f64 to f32

* trans f32 to f16

* fix

* fix

* revert and add strictinfermeta for min/max

* fix

* fix

* fix

* add StrictInfermeta to _INFERMETA_NEED_META_CONFIG

* fix
diff --git a/paddle/fluid/pir/dialect/op_generator/op_build_gen.py b/paddle/fluid/pir/dialect/op_generator/op_build_gen.py
@@ -141,6 +141,8 @@
     'Pool2DInferMeta',
     'ReduceIntArrayAxisInferMetaBase',
     'ReduceIntArrayAxisInferMeta',
+    'StrictReduceIntArrayAxisInferMetaBase',
+    'StrictReduceIntArrayAxisInferMeta',
     'ReshapeInferMeta',
     'ReshapeWithXShapeInferMeta',
     'ReverseInferMeta',
diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
@@ -3713,6 +3713,89 @@ DDim ReduceInferDim(const MetaTensor& x,
   return out_dim;
 }
 
+DDim StrictReduceInferDim(const MetaTensor& x,
+                          const std::vector<int64_t>& axis,
+                          bool keep_dim,
+                          bool reduce_all) {
+  const int x_rank = x.dims().size();
+  uint32_t axis_bitmap = 0;
+
+  for (size_t i = 0; i < axis.size(); ++i) {
+    int64_t formatted_idx = axis[i];
+    if (x_rank == 0) {
+      PADDLE_ENFORCE_EQ(
+          axis[i] == 0 || axis[i] == -1,
+          true,
+          common::errors::InvalidArgument(
+              "When input 0D Tensor, the axis can only be -1, 0, None or []"));
+      formatted_idx = 0;
+    } else {
+      PADDLE_ENFORCE_LT(
+          axis[i],
+          x_rank,
+          errors::InvalidArgument(
+              "The reduce dim index %d should be in the "
+              "range [ -dimension(X), dimension(X) ) "
+              "which dimension = %d. But received dim index = %d.",
+              i,
+              x_rank,
+              axis[i]));
+      PADDLE_ENFORCE_GE(
+          axis[i],
+          -x_rank,
+          errors::InvalidArgument(
+              "The reduce dim index %d should be in the "
+              "range [ -dimension(X), dimension(X) )  "
+              "which dimension = %d. But received dim index = %d.",
+              i,
+              x_rank,
+              axis[i]));
+      if (axis[i] < 0) {
+        formatted_idx += x_rank;
+      }
+    }
+
+    uint32_t bit = 1U << formatted_idx;
+    PADDLE_ENFORCE_EQ(axis_bitmap & bit,
+                      0,
+                      common::errors::InvalidArgument(
+                          "Axis contains duplicate dimensions. Dimension %d "
+                          "appears more than once in axis.",
+                          formatted_idx));
+    axis_bitmap |= bit;
+  }
+
+  bool full_dim = true;
+  if (axis.size() > 0) {
+    uint32_t all_bits = (1U << x_rank) - 1;
+    full_dim = (axis_bitmap == all_bits);
+  }
+  bool empty_dim = axis.size() == 0;
+  reduce_all = reduce_all || full_dim || empty_dim;
+
+  std::vector<int64_t> out_dim_vector;
+  for (int i = 0; i < x_rank; ++i) {
+    uint32_t bit = 1U << i;
+    if (reduce_all || (axis_bitmap & bit)) {
+      PADDLE_ENFORCE_NE(
+          x.dims().at(i),
+          0,
+          common::errors::InvalidArgument(
+              "Cannot perform reduction along an axis (%d) that has a size of "
+              "0. ",
+              i));
+      if (keep_dim) {
+        out_dim_vector.push_back(1);
+      }
+    } else {
+      out_dim_vector.push_back(x.dims().at(i));
+    }
+  }
+
+  DDim out_dim = common::make_ddim(out_dim_vector);
+  return out_dim;
+}
+
 void ReduceInferMetaBase(const MetaTensor& x,
                          const std::vector<int64_t>& axis,
                          bool keep_dim,
@@ -3804,6 +3887,36 @@ void ReduceIntArrayAxisInferMeta(const MetaTensor& x,
   ReduceIntArrayAxisInferMetaBase(x, axis, keep_dim, reduce_all, out, config);
 }
 
+void StrictReduceIntArrayAxisInferMetaBase(const MetaTensor& x,
+                                           const IntArray& axis,
+                                           bool keep_dim,
+                                           bool reduce_all,
+                                           MetaTensor* out,
+                                           MetaConfig config) {
+  DDim out_dim;
+  if (config.is_runtime || !axis.FromTensor()) {
+    out_dim = StrictReduceInferDim(x, axis.GetData(), keep_dim, reduce_all);
+  } else {
+    out_dim = ReduceInferDimForIntArrayAxis(x, axis, keep_dim, reduce_all);
+  }
+  out->set_dims(out_dim);
+  out->set_dtype(x.dtype());
+  out->set_layout(x.layout());
+}
+
+void StrictReduceIntArrayAxisInferMeta(const MetaTensor& x,
+                                       const IntArray& axis,
+                                       bool keep_dim,
+                                       MetaTensor* out,
+                                       MetaConfig config) {
+  bool reduce_all = false;
+  if (axis.size() == 0) {
+    reduce_all = true;
+  }
+  StrictReduceIntArrayAxisInferMetaBase(
+      x, axis, keep_dim, reduce_all, out, config);
+}
+
 void ReduceScatterInferMeta(const MetaTensor& x, int nranks, MetaTensor* out) {
   auto dim = x.dims();
   if (dim[0] > 0 || dim[0] < -1) {
diff --git a/paddle/phi/infermeta/unary.h b/paddle/phi/infermeta/unary.h
@@ -616,6 +616,19 @@ void ReduceIntArrayAxisInferMeta(const MetaTensor& x,
                                  MetaTensor* out,
                                  MetaConfig config = MetaConfig());
 
+void StrictReduceIntArrayAxisInferMetaBase(const MetaTensor& x,
+                                           const IntArray& axis,
+                                           bool keep_dim,
+                                           bool reduce_all,
+                                           MetaTensor* out,
+                                           MetaConfig config = MetaConfig());
+
+void StrictReduceIntArrayAxisInferMeta(const MetaTensor& x,
+                                       const IntArray& axis,
+                                       bool keep_dim,
+                                       MetaTensor* out,
+                                       MetaConfig config = MetaConfig());
+
 void ReduceScatterInferMeta(const MetaTensor& x, int nranks, MetaTensor* out);
 
 void RepeatInterleaveInferMeta(const MetaTensor& x,
diff --git a/paddle/phi/kernels/cpu/reduce_max_kernel.cc b/paddle/phi/kernels/cpu/reduce_max_kernel.cc
@@ -27,6 +27,10 @@ void MaxKernel(const Context& dev_ctx,
                const IntArray& dims,
                bool keep_dim,
                DenseTensor* out) {
+  if (x.numel() == 0) {
+    dev_ctx.template Alloc<T>(out);
+    return;
+  }
   bool reduce_all = recompute_reduce_all(x, dims);
   auto out_dtype = x.dtype();
   phi::Reduce<CPUContext, T, phi::funcs::MaxFunctor>(
diff --git a/paddle/phi/kernels/kps/reduce_kernel.cu b/paddle/phi/kernels/kps/reduce_kernel.cu
@@ -108,6 +108,10 @@ void MaxKernel(const Context& dev_ctx,
                const IntArray& dims,
                bool keep_dim,
                DenseTensor* out) {
+  if (x.numel() == 0) {
+    dev_ctx.template Alloc<T>(out);
+    return;
+  }
   bool reduce_all = recompute_reduce_all(x, dims);
   phi::MaxRawKernel<T, Context>(dev_ctx, x, dims, keep_dim, reduce_all, out);
 }
diff --git a/paddle/phi/kernels/onednn/reduce_max_kernel.cc b/paddle/phi/kernels/onednn/reduce_max_kernel.cc
@@ -25,6 +25,10 @@ void MaxKernel(const Context& dev_ctx,
                const IntArray& dims,
                bool keep_dim,
                DenseTensor* out) {
+  if (x.numel() == 0) {
+    dev_ctx.template Alloc<T>(out);
+    return;
+  }
   bool reduce_all = recompute_reduce_all(x, dims);
   ReduceKernel<T, Context>(dev_ctx,
                            x,
diff --git a/paddle/phi/kernels/reduce_min_kernel.cc b/paddle/phi/kernels/reduce_min_kernel.cc
@@ -26,12 +26,11 @@ void MinKernel(const Context& dev_ctx,
                const IntArray& dims,
                bool keep_dim,
                DenseTensor* out) {
+  if (x.numel() == 0) {
+    dev_ctx.template Alloc<T>(out);
+    return;
+  }
   bool reduce_all = recompute_reduce_all(x, dims);
-  PADDLE_ENFORCE_GT(
-      x.numel(),
-      0,
-      errors::InvalidArgument("Zero-size tensor to reduction operation minimum "
-                              "which has no identity."));
   MinRawKernel<T>(dev_ctx, x, dims, keep_dim, reduce_all, out);
 }
 
diff --git a/paddle/phi/kernels/xpu/reduce_max_kernel.cc b/paddle/phi/kernels/xpu/reduce_max_kernel.cc
@@ -27,6 +27,10 @@ void MaxKernel(const Context& dev_ctx,
                const IntArray& dims,
                bool keep_dim,
                DenseTensor* out) {
+  if (x.numel() == 0) {
+    dev_ctx.template Alloc<T>(out);
+    return;
+  }
   bool reduce_all = recompute_reduce_all(x, dims);
   using XPUType = typename XPUTypeTrait<T>::Type;
   auto f = [](xpu::Context* ctx,
diff --git a/paddle/phi/ops/yaml/inconsistent/dygraph_ops.yaml b/paddle/phi/ops/yaml/inconsistent/dygraph_ops.yaml
@@ -257,7 +257,7 @@
   args : (Tensor x, IntArray axis={}, bool keepdim=false)
   output : Tensor(out)
   infer_meta :
-    func : ReduceIntArrayAxisInferMeta
+    func : StrictReduceIntArrayAxisInferMeta
     spmd_rule : ReductionMinInferSpmdDynamic
   kernel :
     func : min
diff --git a/paddle/phi/ops/yaml/inconsistent/static_ops.yaml b/paddle/phi/ops/yaml/inconsistent/static_ops.yaml
@@ -520,7 +520,7 @@
   args : (Tensor x, IntArray axis={}, bool keepdim=false)
   output : Tensor(out)
   infer_meta :
-    func : ReduceIntArrayAxisInferMeta
+    func : StrictReduceIntArrayAxisInferMeta
     spmd_rule : ReductionMinInferSpmdDynamic
   kernel :
     func : min
diff --git a/paddle/phi/ops/yaml/legacy/static_ops.yaml b/paddle/phi/ops/yaml/legacy/static_ops.yaml
@@ -549,7 +549,7 @@
   args : (Tensor x, IntArray axis={0}, bool keepdim=false, bool reduce_all=false, int in_dtype=-1, int out_dtype=-1)
   output : Tensor(out)
   infer_meta :
-    func : ReduceIntArrayAxisInferMetaBase
+    func : StrictReduceIntArrayAxisInferMetaBase
     param : [x, axis, keepdim, reduce_all]
   kernel :
     func : max_raw
@@ -569,7 +569,7 @@
   args : (Tensor x, IntArray axis={0}, bool keepdim=false, bool reduce_all=false, int in_dtype=-1, int out_dtype=-1)
   output : Tensor(out)
   infer_meta :
-    func : ReduceIntArrayAxisInferMetaBase
+    func : StrictReduceIntArrayAxisInferMetaBase
     param : [x, axis, keepdim, reduce_all]
   kernel :
     func : min_raw
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -3439,7 +3439,7 @@
   args : (Tensor x, IntArray axis={}, bool keepdim=false)
   output : Tensor(out)
   infer_meta :
-    func : ReduceIntArrayAxisInferMeta
+    func : StrictReduceIntArrayAxisInferMeta
     spmd_rule: ReductionMaxInferSpmdDynamic
   kernel :
     func : max
diff --git a/test/legacy_test/test_max_op.py b/test/legacy_test/test_max_op.py
@@ -22,6 +22,7 @@
 sys.path.append("../../legacy_test")
 import os
 
+from op_test import OpTest
 from test_sum_op import TestReduceOPTensorAxisBase
 from utils import dygraph_guard, static_guard
 
@@ -145,6 +146,113 @@ def init_data(self):
         ]
 
 
+class TestMaxZeroSize1(unittest.TestCase):
+    def init_data(self):
+        self.shape = [0, 1, 2, 3]
+        self.axis = [1, 2, 3]
+        self.keepdims = False
+
+    def setUp(self):
+        self.init_data()
+        self.data = np.random.random(self.shape).astype(np.float64)
+        self.expect_res = np.max(
+            self.data, axis=tuple(self.axis), keepdims=self.keepdims
+        )
+        self.places = [core.CPUPlace()]
+        if core.is_compiled_with_cuda():
+            self.places.append(core.CUDAPlace(0))
+
+    def test_static(self):
+        with static_guard():
+            for place in self.places:
+                with paddle.static.program_guard(
+                    paddle.static.Program(), paddle.static.Program()
+                ):
+                    x = paddle.static.data(
+                        "x", shape=self.shape, dtype="float64"
+                    )
+                    res = paddle.max(x, axis=self.axis, keepdim=self.keepdims)
+                    exe = paddle.static.Executor(place)
+                    (res,) = exe.run(feed={"x": self.data}, fetch_list=[res])
+                np.testing.assert_equal(res, self.expect_res)
+
+    def test_dygraph(self):
+        with dygraph_guard():
+            x = paddle.to_tensor(self.data)
+            res = paddle.max(x, axis=self.axis, keepdim=self.keepdims)
+        np.testing.assert_equal(res, self.expect_res)
+
+
+class TestMaxZeroSize2(TestMaxZeroSize1):
+    def init_data(self):
+        self.shape = [0, 0, 2]
+        self.axis = [2]
+        self.keepdims = False
+
+
+class TestMaxZeroSize3(TestMaxZeroSize1):
+    def init_data(self):
+        self.shape = [0, 0, 2]
+        self.axis = [2]
+        self.keepdims = True
+
+
+class TestMaxOp(OpTest):
+    def setUp(self):
+        self.op_type = "reduce_max"
+        self.python_api = paddle.max
+        self.init_data()
+        self.prepare_data()
+
+    def init_data(self):
+        self.shape = [0, 1, 2]
+        self.axis = [1]
+        self.dtype = np.float64
+        self.keepdims = False
+
+    def prepare_data(self):
+        self._input_data = np.random.random(self.shape).astype(self.dtype)
+        self._output_data = np.max(
+            self._input_data, keepdims=self.keepdims, axis=tuple(self.axis)
+        )
+        self.inputs = {'X': self._input_data}
+        self.outputs = {'Out': self._output_data}
+        self.attrs = {"dim": self.axis, "keep_dim": self.keepdims}
+
+    def test_check_output(self):
+        self.check_output(check_pir=True)
+
+    def test_check_grad(self):
+        self.check_grad(
+            ['X'],
+            ['Out'],
+            check_pir=True,
+        )
+
+
+@unittest.skipIf(
+    not core.supports_bfloat16(), "place does not support BF16 evaluation"
+)
+class TestMaxBfloat16(unittest.TestCase):
+    def init_data(self):
+        self.shape = [0, 1, 2]
+        self.axis = [1]
+        self.keepdims = False
+
+    def setUp(self):
+        self.init_data()
+        data = np.random.random(self.shape).astype(np.float64)
+        res = np.max(data, axis=tuple(self.axis), keepdims=self.keepdims)
+        self.expect_shape = res.shape
+
+    def test_shape(self):
+        with dygraph_guard():
+            x = paddle.zeros(self.shape, dtype=paddle.bfloat16)
+            res = paddle.max(x, axis=self.axis, keepdim=self.keepdims)
+            res = res.numpy()
+            np.testing.assert_equal(res.shape, self.expect_shape)
+
+
 class TestMaxWithNan(unittest.TestCase):
     def _get_places(self):
         places = []
diff --git a/test/legacy_test/test_min_op.py b/test/legacy_test/test_min_op.py
