diff --git a/onnxruntime/core/providers/cpu/tensor/scatter.cc b/onnxruntime/core/providers/cpu/tensor/scatter.cc index 230deaecd2..be6e19b80e 100644 --- a/onnxruntime/core/providers/cpu/tensor/scatter.cc +++ b/onnxruntime/core/providers/cpu/tensor/scatter.cc @@ -106,7 +106,7 @@ Status CopyScatterData(const Tensor* data_input, const Tensor* indices_input, co // We start at num_dims - 2 because we already pre-populated // the last element above for (int64_t i = int64_t(num_dims - 2); i >= 0; --i) { - dim_block_size[i] = input_data_shape[i] * dim_block_size[i + 1]; + dim_block_size[i] = input_data_shape[i + 1] * dim_block_size[i + 1]; } } diff --git a/onnxruntime/test/providers/cpu/tensor/scatter_op_test.cc b/onnxruntime/test/providers/cpu/tensor/scatter_op_test.cc index 57d07d389d..fc9d72a07d 100644 --- a/onnxruntime/test/providers/cpu/tensor/scatter_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/scatter_op_test.cc @@ -124,5 +124,27 @@ TEST(ScatterOpTest, IndicesUpdatesDimsDonotMatch) { test.AddOutput("y", {1, 5}, {1.0f, 1.1f, 3.0f, 2.1f, 5.0f}); test.Run(OpTester::ExpectResult::kExpectFailure, "Indices vs updates dimensions differs at position=1 3 vs 2"); } + +TEST(ScatterOpTest, ValidIndex) { + OpTester test("Scatter", Scatter_ver); + test.AddAttribute("axis", 0); + + test.AddInput("data", {4, 2, 1}, {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}); + test.AddInput("indices", {1, 1, 1}, {3}); + test.AddInput("updates", {1, 1, 1}, {5.0f}); + test.AddOutput("y", {4, 2, 1}, {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 5.0f, 0.0f}); + test.Run(); +} + +TEST(ScatterOpTest, InvalidIndex) { + OpTester test("Scatter", Scatter_ver); + test.AddAttribute("axis", 0); + + test.AddInput("data", {4, 2, 1}, {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}); + test.AddInput("indices", {1, 1, 1}, {4}); + test.AddInput("updates", {1, 1, 1}, {5.0f}); + test.AddOutput("y", {4, 2, 1}, {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 5.0f, 0.0f}); + test.Run(OpTester::ExpectResult::kExpectFailure, "indices element out of data bounds, idx=4 data_dim=4"); +} } // namespace test } // namespace onnxruntime