【Infer Symbolic Shape BUAA No.76】Add mean_all op (#66609)

crazyxiaoxi · web-flow · commit ad84e3a09e3c · 2024-08-01T10:02:47.000+08:00
* mean_all

* test_op

* nothing

* test mean_all

* test mean

* test mean

* add mean_all optest
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc
@@ -637,6 +637,28 @@ bool MinOpInferSymbolicShape(pir::Operation *op,
   return MaxOpInferSymbolicShape(op, infer_context);
 }
 
+bool MeanAllOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const std::vector<symbol::DimExpr> &x_dims = x_shape_or_data.shape();
+
+  PADDLE_ENFORCE_GT(
+      x_dims.size(),
+      0,
+      phi::errors::InvalidArgument("Input(x) of MeanAllOp must have rank "
+                                   "greater than 0, but received rank 0."));
+
+  std::vector<symbol::DimExpr> output_shape = {};
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(output_shape)});
+
+  return true;
+}
+
 bool NonzeroOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
   const auto &x_shape_or_data =
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.h
@@ -50,6 +50,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logsumexp)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Max)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Maxout)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Min)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(MeanAll)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Nonzero)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Numel)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Pad)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -3086,6 +3086,7 @@
   kernel :
     func : mean_all
   backward : mean_all_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : memcpy_d2h
   args : (Tensor x, int dst_place_type)
diff --git a/test/legacy_test/test_mean_op.py b/test/legacy_test/test_mean_op.py
@@ -65,6 +65,30 @@ def test_checkout_grad(self):
         self.check_grad(['X'], 'Out', check_pir=True, check_prim_pir=True)
 
 
+class TestMeanAllOp(OpTest):
+    def setUp(self):
+        self.op_type = "mean_all"
+        self.python_api = paddle.mean_all
+        self.public_python_api = paddle.mean_all
+        self.dtype = np.float64
+        self.init_dtype_type()
+        self.init_prim_type()
+        self.inputs = {'X': np.random.random((10, 10)).astype(self.dtype)}
+        self.outputs = {'Out': np.mean(self.inputs["X"])}
+
+    def init_prim_type(self):
+        self.prim_op_type = "comp"
+
+    def init_dtype_type(self):
+        pass
+
+    def test_check_output(self):
+        self.check_output(check_pir=True)
+
+    def test_checkout_grad(self):
+        self.check_grad(['X'], 'Out', check_pir=True, check_prim_pir=True)
+
+
 class TestMeanOpPrim(TestMeanOp):
     def init_prim_type(self):
         self.prim_op_type = "prim"
@@ -90,6 +114,26 @@ def test_checkout_grad(self):
         self.check_grad(['X'], 'Out', check_pir=True, check_prim_pir=True)
 
 
+class TestMeanAllOp_ZeroDim(OpTest):
+    def setUp(self):
+        self.op_type = "mean_all"
+        self.python_api = paddle.mean_all
+        self.dtype = np.float64
+        self.public_python_api = paddle.mean_all
+        self.init_prim_type()
+        self.inputs = {'X': np.random.random([]).astype(self.dtype)}
+        self.outputs = {'Out': np.mean(self.inputs["X"])}
+
+    def init_prim_type(self):
+        self.prim_op_type = "comp"
+
+    def test_check_output(self):
+        self.check_output(check_pir=True)
+
+    def test_checkout_grad(self):
+        self.check_grad(['X'], 'Out', check_pir=True, check_prim_pir=True)
+
+
 class TestMeanOp_ZeroDim_Prim(TestMeanOp_ZeroDim):
     def init_prim_type(self):
         self.prim_op_type = "prim"
