[Infer Symbolic Shape No.26][BUAA]dropout (#67412)

MufanColin · web-flow · commit 288c616fa8db · 2024-08-15T14:00:19.000+08:00
* Fixed dropout op

* Removed typos

* Tried a new version

* Added check_symbol_infer=False

* removed self.check_output
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
@@ -367,6 +367,22 @@ bool CrossOpInferSymbolicShape(pir::Operation *op,
 //   return true;
 // }
 
+bool DropoutOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const symbol::ShapeOrDataDimExprs &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+
+  infer_context->SetShapeOrDataForValue(op->result(0), x_shape_or_data);
+
+  if (op->result(1)) {
+    symbol::TensorShapeOrDataDimExprs mask_shape(x_shape_or_data.shape());
+    infer_context->SetShapeOrDataForValue(
+        op->result(1), symbol::ShapeOrDataDimExprs{mask_shape});
+  }
+
+  return true;
+}
+
 bool EmbeddingOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
   const std::vector<symbol::DimExpr> &x_dims =
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.h
@@ -33,6 +33,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Conv2d)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Conv3d)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cross)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Dot)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Dropout)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Embedding)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(EqualAll)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(SparseWeightEmbedding)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -1464,6 +1464,7 @@
   intermediate : mask
   backward : dropout_grad
   traits : pir::SideEffectTrait
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : edit_distance
   args : (Tensor hyps, Tensor refs, Tensor hypslength, Tensor refslength, bool normalized = false)
