[buaa]【Infer Symbolic Shape No.63】fix index_select_strided & margin op (#67115)

crazyxiaoxi · web-flow · commit 01c98f9288e6 · 2024-08-14T10:37:01.000+08:00
* test_index_select_strided

* testop

* delete error annotation

* annotation
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
@@ -875,7 +875,7 @@ bool MarginCrossEntropyOpInferSymbolicShape(
   const auto &one = symbol::DimExpr{1};
 
   if (labels_dims.size() > 1) {
-    infer_context->AddEqualCstr(labels_dims[axis - 1], one);
+    infer_context->AddEqualCstr(labels_dims[axis], one);
   }
 
   infer_context->SetShapeOrDataForValue(
@@ -1048,13 +1048,13 @@ bool IndexSelectStridedOpInferSymbolicShape(
   std::vector<symbol::DimExpr> output_dims(input_dims.begin(),
                                            input_dims.end());
   output_dims.erase(output_dims.begin() + dim);
+  // No need to add any constraints here as we are simply removing a dimension.
 
   infer_context->SetShapeOrDataForValue(
       op->result(0),
       symbol::ShapeOrDataDimExprs{
           symbol::TensorShapeOrDataDimExprs(output_dims)});
 
-  // No need to add any constraints here as we are simply removing a dimension.
   return true;
 }
 
