diff --git a/mindspore/ccsrc/pipeline/jit/static_analysis/prim.cc b/mindspore/ccsrc/pipeline/jit/static_analysis/prim.cc index fae6f032286..01d66244f0e 100644 --- a/mindspore/ccsrc/pipeline/jit/static_analysis/prim.cc +++ b/mindspore/ccsrc/pipeline/jit/static_analysis/prim.cc @@ -90,6 +90,16 @@ mindspore::HashSet prims_use_sequence_elements{ prim::kPrimSparseTensorDenseMatmul->name(), prim::kPrimBroadcast->name(), prim::kPrimEinsumGrad->name(), + prim::kPrimFlattenGrad->name(), + prim::kPrimMirror->name(), + prim::kPrimMirrorMiniStep->name(), + prim::kPrimMiniStepAllGather->name(), + prim::kPrimMicroStepAllGather->name(), + prim::kPrimVirtualDiv->name(), + prim::kPrimVirtualAdd->name(), + prim::kPrimVirtualDataset->name(), + prim::kPrimVirtualOutput->name(), + prim::kPrimFill->name(), "InvertPermutation", "Meshgrid", "TransShape", diff --git a/mindspore/core/abstract/prim_arrays.cc b/mindspore/core/abstract/prim_arrays.cc index 1bb34566163..f8f6e669175 100644 --- a/mindspore/core/abstract/prim_arrays.cc +++ b/mindspore/core/abstract/prim_arrays.cc @@ -1369,6 +1369,8 @@ AbstractBasePtr InferImplDynamicStitch(const AnalysisEnginePtr &, const Primitiv min_shape[0] = 1; max_shape[0] = indices_total_size * EXPAND_MAX; } + SetSequenceElementsUseFlags(input_tuple, true); + SetSequenceElementsUseFlags(input_tuple_1, true); return std::make_shared(infer_type, std::make_shared(out_shape, min_shape, max_shape)); } @@ -1379,6 +1381,9 @@ AbstractBasePtr InferImplTensorCopySlices(const AnalysisEnginePtr &, const Primi constexpr auto kTensorCopySlicesInputNum = 5; CheckArgsSize(op_name, args_spec_list, kTensorCopySlicesInputNum); AbstractTensorPtr input = CheckArg(op_name, args_spec_list, 0); + for (size_t i = 2; i < args_spec_list.size(); ++i) { + SetSequenceElementsUseFlags(args_spec_list[i], true); + } return std::make_shared(input->element(), input->shape()); } } // namespace abstract