【Infer Symbolic Shape BUAA No.10 No.73 No.78】Add mode , margin_cross_entropy , box_clip. op

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

@@ -720,6 +720,46 @@ bool MatmulOpInferSymbolicShape(pir::Operation *op,
   return true;
 }
 
+bool MarginCrossEntropyOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &logits_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const auto &labels_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1));
+
+  std::vector<symbol::DimExpr> logits_dims = logits_shape_or_data.shape();
+  std::vector<symbol::DimExpr> labels_dims = labels_shape_or_data.shape();
+
+  size_t logits_rank = logits_dims.size();
+  auto axis = logits_rank - 1;
+
+  for (size_t i = 0; i < logits_rank; i++) {
+    if (i != axis) {
+      infer_context->AddBroadcastableCstr(logits_dims[i], labels_dims[i]);
+    }
+  }
+
+  const auto one = symbol::DimExpr{1};
+
+  if (labels_dims.size() > 1) {
+    infer_context->AddEqualCstr(logits_dims[axis], one);
+  }
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(logits_dims)});
+
+  logits_dims[axis] = symbol::DimExpr(1);
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(1),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(logits_dims)});
+
+  return true;
+}
+
 bool MvOpInferSymbolicShape(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

@@ -52,6 +52,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Kron)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(MatrixRankTol)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(MaskedSelect)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Matmul)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(MarginCrossEntropy)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Mv)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Nextafter)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(PullBoxSparse)

paddle/phi/ops/yaml/ops.yaml

@@ -3006,6 +3006,7 @@
     func : margin_cross_entropy
     data_type : logits
   backward : margin_cross_entropy_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : masked_multihead_attention_
   args : (Tensor x, Tensor cache_kv, Tensor bias, Tensor src_mask, Tensor cum_offsets, Tensor sequence_lengths, Tensor rotary_tensor, Tensor beam_cache_offset, Tensor qkv_out_scale, Tensor out_shift, Tensor out_smooth, int seq_len, int rotary_emb_dims, bool use_neox_rotary_style=false, str compute_dtype = "default", float out_scale=-1, int quant_round_type=1, float quant_max_bound=127.0, float quant_min_bound=-127.0)