diff --git a/onnxruntime/core/providers/cpu/controlflow/scan_9.cc b/onnxruntime/core/providers/cpu/controlflow/scan_9.cc index 7dffa6cef4..b0dbbd36e4 100644 --- a/onnxruntime/core/providers/cpu/controlflow/scan_9.cc +++ b/onnxruntime/core/providers/cpu/controlflow/scan_9.cc @@ -343,7 +343,7 @@ Status ScanImpl::SetupInputs() { std::vector permutations; std::vector new_shape; - CalculateTransposedShape(input_shape, sequence_dim, permutations, new_shape); + CalculateTransposedShapeForInput(input_shape, sequence_dim, permutations, new_shape); if (!alloc) { status = context_.GetTempSpaceAllocator(&alloc); @@ -474,7 +474,7 @@ Status ScanImpl::TransposeOutput() { std::vector permutations; std::vector new_shape; - CalculateTransposedShape(temporary_output_tensor.Shape(), axis, permutations, new_shape); + CalculateTransposedShapeForOutput(temporary_output_tensor.Shape(), axis, permutations, new_shape); Tensor* output = context_.Output(output_index, new_shape); ORT_ENFORCE(output, "Outputs from Scan are not optional and should never be null."); diff --git a/onnxruntime/core/providers/cpu/controlflow/scan_utils.cc b/onnxruntime/core/providers/cpu/controlflow/scan_utils.cc index d02cd116d6..c3e7fbcfee 100644 --- a/onnxruntime/core/providers/cpu/controlflow/scan_utils.cc +++ b/onnxruntime/core/providers/cpu/controlflow/scan_utils.cc @@ -211,25 +211,47 @@ MLValue AllocateTensorInMLValue(const MLDataType data_type, const TensorShape& s DataTypeImpl::GetType()->GetDeleteFunc()}; }; -void CalculateTransposedShape(const TensorShape& input_shape, int64_t axis, - std::vector& permutations, std::vector& output_shape) { - int64_t rank = input_shape.NumDimensions(); - const auto& dims = input_shape.GetDims(); +void CalculateTransposedShapeForInput(const TensorShape& original_shape, int64_t axis, + std::vector& permutations, std::vector& transposed_shape) { + int64_t rank = original_shape.NumDimensions(); + const auto& dims = original_shape.GetDims(); permutations.reserve(rank); permutations.push_back(axis); - output_shape.reserve(rank); - output_shape.push_back(dims[axis]); + transposed_shape.reserve(rank); + transposed_shape.push_back(dims[axis]); for (int64_t i = 0; i < rank; ++i) { if (i != axis) { permutations.push_back(i); - output_shape.push_back(dims[i]); + transposed_shape.push_back(dims[i]); } } } +void CalculateTransposedShapeForOutput(const TensorShape& original_shape, int64_t axis, + std::vector& permutations, std::vector& transposed_shape) { + int64_t rank = original_shape.NumDimensions(); + const auto& dims = original_shape.GetDims(); + + permutations.reserve(rank); + transposed_shape.reserve(rank); + + for (int64_t i = 1; i <= axis; ++i) { + permutations.push_back(i); + transposed_shape.push_back(dims[i]); + } + + permutations.push_back(0); + transposed_shape.push_back(dims[0]); + + for (int64_t i = axis + 1; i < rank; ++i) { + permutations.push_back(i); + transposed_shape.push_back(dims[i]); + } +} + LoopStateVariable::LoopStateVariable(const MLValue& original_value, MLValue& final_value, const int64_t sequence_len, diff --git a/onnxruntime/core/providers/cpu/controlflow/scan_utils.h b/onnxruntime/core/providers/cpu/controlflow/scan_utils.h index eef4dc971e..c61277b85a 100644 --- a/onnxruntime/core/providers/cpu/controlflow/scan_utils.h +++ b/onnxruntime/core/providers/cpu/controlflow/scan_utils.h @@ -173,14 +173,23 @@ Status IterateSequence(OpKernelContextInternal& context, MLValue AllocateTensorInMLValue(const MLDataType data_type, const TensorShape& shape, AllocatorPtr& allocator); /** -Calculate the transpose permutations and output shape by shifting the chosen axis to the first dimension. +Calculate the transpose permutations and shape by shifting the chosen axis TO the first dimension. The other dimension indexes or values are pushed in order after the chosen axis. e.g. if shape is {2, 3, 4} and axis 1 is chosen the permutations will be {1, 0, 2} and output shape will be {3, 2, 4} if axis 2 is chosen the permutations will be {2, 0, 1} and the output shape will be {4, 2, 3} */ -void CalculateTransposedShape(const TensorShape& input_shape, int64_t axis, - std::vector& permutations, std::vector& output_shape); +void CalculateTransposedShapeForInput(const TensorShape& original_shape, int64_t axis, + std::vector& permutations, std::vector& transposed_shape); + +/** +Calculate the transpose permutations and shape by shifting the chosen axis FROM the first dimension. + +e.g. if shape is {4, 2, 3} and axis 2 is chosen, dimension 0 will move to dimension 2, + the permutations will be {1, 2, 0} and output shape will be {2, 3, 4} +*/ +void CalculateTransposedShapeForOutput(const TensorShape& original_shape, int64_t axis, + std::vector& permutations, std::vector& transposed_shape); } // namespace detail } // namespace scan diff --git a/onnxruntime/test/providers/cpu/controlflow/scan_test.cc b/onnxruntime/test/providers/cpu/controlflow/scan_test.cc index d140a3f560..1894698335 100644 --- a/onnxruntime/test/providers/cpu/controlflow/scan_test.cc +++ b/onnxruntime/test/providers/cpu/controlflow/scan_test.cc @@ -382,8 +382,8 @@ static void RunTest_v9(const std::string test_name, int64_t sequence_len, int64_ if (axis >= -rank && axis < rank) { std::vector permutations; std::vector new_shape; - scan::detail::CalculateTransposedShape(output_shape, HandleNegativeAxis(axis, output_shape.size()), - permutations, new_shape); + scan::detail::CalculateTransposedShapeForOutput(output_shape, HandleNegativeAxis(axis, output_shape.size()), + permutations, new_shape); return new_shape; } } @@ -844,6 +844,48 @@ TEST(Scan9, TransposeOutput) { iteration_count_out, output_0, output_1, output_2, output_3); } +TEST(Scan9, TransposeOutputDim2) { + // Construct scan body subgraph with 1 scan inputs, 1 scan outputs + // scan-in-1 => scan-out-1 + Model model("ScanBody"); + auto& graph = model.MainGraph(); + + TypeProto float_tensor; + float_tensor.mutable_tensor_type()->set_elem_type(TensorProto_DataType_FLOAT); + float_tensor.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); + float_tensor.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); + + auto& scan_in_1 = graph.GetOrCreateNodeArg("scan_in_1", &float_tensor); + auto& scan_out_1 = graph.GetOrCreateNodeArg("scan_out_1", &float_tensor); + + graph.AddNode("pass_through", "Identity", "Copy scan_in_1 to scan_out_1", {&scan_in_1}, {&scan_out_1}); + + auto status = graph.Resolve(); + EXPECT_EQ(status, Status::OK()); + + auto& scan_body = graph.ToGraphProto(); + + ScanOpTester test{9}; + + std::vector input_shape{2, 1, 1}; + + // transpose on axis 2, so dim 0 of the output (copied directly from input of {2, 1, 1}) + // will move to dim 2 of the output giving shape {1, 1, 2} + std::vector output_axes{2}; + std::vector output_shape{1, 1, 2}; + + test.AddAttribute("body", scan_body); + test.AddAttribute("num_scan_inputs", 1); + test.AddAttribute>("scan_output_axes", output_axes); + + // the data won't change, but the shape should be transposed from 2, 1, 1 to 1, 1, 2, which + // OpTester::Run will validate + test.AddInput("scan_input_1", input_shape, {1.0, 2.0}); + test.AddOutput("scan_output_1", output_shape, {1.0, 2.0}); + + test.Run(); +} + static void InvalidInput(bool is_v8) { const int64_t batch_size = 1; const int64_t sequence_len = 2;