Merge pull request #49815 from geetachavan1/cherrypicks_X3FU2

Prevent array write out-of-bounds.
tensorflow · May 27, 2021 · c427d9b · c427d9b
2 parents a2c5283 + 904178c
commit c427d9b
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tensorflow/lite/kernels/arg_min_max.cc b/tensorflow/lite/kernels/arg_min_max.cc
@@ -36,6 +36,9 @@ TfLiteStatus ResizeOutput(TfLiteContext* context, const TfLiteTensor* input,
     axis_value += NumDimensions(input);
   }
 
+  TF_LITE_ENSURE(context, axis_value >= 0);
+  TF_LITE_ENSURE(context, axis_value < NumDimensions(input));
+
   // Copy the input dimensions to output except the axis dimension.
   TfLiteIntArray* output_dims = TfLiteIntArrayCreate(NumDimensions(input) - 1);
   int j = 0;