diff --git a/onnxruntime/core/providers/cpu/ml/array_feature_extractor.cc b/onnxruntime/core/providers/cpu/ml/array_feature_extractor.cc index af67419f4fb91..60ebf862e1601 100644 --- a/onnxruntime/core/providers/cpu/ml/array_feature_extractor.cc +++ b/onnxruntime/core/providers/cpu/ml/array_feature_extractor.cc @@ -73,10 +73,10 @@ common::Status ArrayFeatureExtractorOp::Compute(OpKernelContext* context) con } for (int64_t i = 0; i < num_indices; ++i) { - if (y_data[i] >= stride) { + if (y_data[i] < 0 || y_data[i] >= stride) { return ORT_MAKE_STATUS( ONNXRUNTIME, INVALID_ARGUMENT, - "Invalid Y argument: index is out of range: Y[", i, "] (", y_data[i], ") >=", stride); + "Invalid Y argument: index is out of range: Y[", i, "] (", y_data[i], ") must be in [0, ", stride, ")"); } } diff --git a/onnxruntime/test/providers/cpu/ml/array_feature_extractor_test.cc b/onnxruntime/test/providers/cpu/ml/array_feature_extractor_test.cc index c7fc73456dcba..671ada7d36383 100644 --- a/onnxruntime/test/providers/cpu/ml/array_feature_extractor_test.cc +++ b/onnxruntime/test/providers/cpu/ml/array_feature_extractor_test.cc @@ -109,5 +109,13 @@ TEST_F(ArrayFeatureExtractorTest, InvalidInputOutOfBoundsY) { test_.Run(OpTester::ExpectResult::kExpectFailure); } +TEST_F(ArrayFeatureExtractorTest, InvalidInputNegativeY) { + test_.AddInput("X", {10, 1}, {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); + test_.AddInput("Y", {1}, {-10}); + // Should fail due to negative index -10 + test_.AddOutput("Z", {0}, {}); + test_.Run(OpTester::ExpectResult::kExpectFailure); +} + } // namespace test } // namespace onnxruntime