[Infer Symbolic Shape No.27,71][BUAA]edit_distance,kldiv_loss (#67327)

* Finished kldiv loss op

* Finished kldiv loss op

* Fixed name issues

* Resolved comments in #67117

Author: MufanColin

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc

@@ -643,6 +643,51 @@ bool IndexSampleOpInferSymbolicShape(
   return true;
 }
 
+bool KldivLossOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const auto &label_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1));
+  const auto &x_shape = x_shape_or_data.shape();
+  const auto &label_shape = label_shape_or_data.shape();
+
+  PADDLE_ENFORCE_EQ(x_shape.size(),
+                    label_shape.size(),
+                    common::errors::InvalidArgument(
+                        "Input(X) rank and Input(Target) rank should be same, "
+                        "but received X rank(%d) != Target rank(%d)",
+                        x_shape.size(),
+                        label_shape.size()));
+
+  for (size_t i = 0; i < x_shape.size(); ++i) {
+    infer_context->AddEqualCstr(x_shape[i], label_shape[i]);
+  }
+
+  std::string reduction =
+      op->attribute<pir::StrAttribute>("reduction").AsString();
+  bool reduction_valid = (reduction == "mean" || reduction == "sum" ||
+                          reduction == "batchmean" || reduction == "none");
+  PADDLE_ENFORCE_EQ(
+      reduction_valid,
+      true,
+      common::errors::InvalidArgument(
+          "Attr(reduction) can only be 'none'|'batchmean'|'sum'|'mean'."));
+
+  std::vector<symbol::DimExpr> out_shape;
+  if (reduction == "none") {
+    out_shape = x_shape;
+  } else {
+    out_shape = std::vector<symbol::DimExpr>{};
+  }
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(out_shape)});
+
+  return true;
+}
+
 bool KronOpInferSymbolicShape(pir::Operation *op,
                               pir::InferSymbolicShapeContext *infer_context) {
   const auto &x_shape_or_data =

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.h

@@ -49,6 +49,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Isclose)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(AccuracyCheck)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(IndexSample)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(IndexSelectStrided)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(KldivLoss)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Kron)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Lstsq)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(MatrixRankTol)

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc

@@ -859,11 +859,10 @@ bool EditDistanceOpInferSymbolicShape(
     infer_context->AddEqualCstr(refs_dims[1], one);
   }
 
-  symbol::ShapeOrDataDimExprs refs_shape_or_data_exprs(
+  symbol::ShapeOrDataDimExprs out_shape_or_data_exprs(
       symbol::TensorShapeOrDataDimExprs(
           std::vector<symbol::DimExpr>{refs_dims}));
-  infer_context->SetShapeOrDataForValue(op->result(0),
-                                        refs_shape_or_data_exprs);
+  infer_context->SetShapeOrDataForValue(op->result(0), out_shape_or_data_exprs);
 
   symbol::ShapeOrDataDimExprs single_dim_expr(symbol::TensorShapeOrDataDimExprs(
       std::vector<symbol::DimExpr>{symbol::DimExpr(1)}));

paddle/phi/ops/yaml/ops.yaml

@@ -2624,6 +2624,7 @@
     func : kldiv_loss
     data_type : x
   backward : kldiv_loss_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : kron
   args : (Tensor x, Tensor y)