add index sample infer symbol (#64332)

phlrain · web-flow · commit 1e8f32d2ecd8 · 2024-05-17T09:41:12.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
@@ -339,6 +339,14 @@ bool GatherNdOpInferSymbolicShape(
   return true;
 }
 
+bool IndexSampleOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const symbol::ShapeOrDataDimExprs &operand_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1));
+  infer_context->SetShapeOrDataForValue(op->result(0), operand_shape_or_data);
+  return true;
+}
+
 bool KronOpInferSymbolicShape(pir::Operation *op,
                               pir::InferSymbolicShapeContext *infer_context) {
   const auto &x_shape_or_data =
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
@@ -26,6 +26,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(ExpandAs)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Gather)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(GatherNd)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Isclose)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(IndexSample)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Kron)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(MaskedSelect)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Matmul)
diff --git a/paddle/phi/api/yaml/ops.yaml b/paddle/phi/api/yaml/ops.yaml
@@ -1630,6 +1630,7 @@
   backward : index_sample_grad
   data_transform :
     skip_transform : index
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : index_select
   args : (Tensor x, Tensor index, int axis = 0)
