【Infer Symbolic Shape No.5】Add infer_symbol_shape for ops binomial (#66595)

enkilee · web-flow · commit f33b758cea1f · 2024-08-09T10:23:47.000+08:00
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
@@ -145,6 +145,34 @@ bool BceLoss_OpInferSymbolicShape(
   return BceLossOpInferSymbolicShape(op, infer_context);
 }
 
+bool BinomialOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &count_shape =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0)).shape();
+  const auto &prob_shape =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1)).shape();
+  PADDLE_ENFORCE_EQ(count_shape.size(),
+                    prob_shape.size(),
+                    common::errors::PreconditionNotMet(
+                        "Input(count) and Input(prob) must have the same "
+                        "dimension size. but got %d vs %d",
+                        count_shape.size(),
+                        prob_shape.size()));
+  for (size_t i = 0; i < count_shape.size(); ++i) {
+    infer_context->AddEqualCstr(count_shape[i], prob_shape[i]);
+  }
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(count_shape)});
+  return true;
+}
+
+bool Binomial_OpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  return BinomialOpInferSymbolicShape(op, infer_context);
+}
+
 // bool BincountOpInferSymbolicShape(pir::Operation *op,
 //                                   pir::InferSymbolicShapeContext
 //                                   *infer_context) {
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
@@ -22,6 +22,8 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Allclose)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Atan2)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(BceLoss)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(BceLoss_)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Binomial)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Binomial_)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Bincount)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Bmm)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(CholeskySolve)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -553,6 +553,7 @@
     func : BinomialInferMeta
   kernel :
     func : binomial
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : bipartite_match
   args: (Tensor dist_mat, str match_type = "bipartite", float dist_threshold = 0.5)
