Merge pull request #49821 from geetachavan1/cherrypicks_N358P

Prevent array OOB read/write
tensorflow · May 27, 2021 · d8ba32d · d8ba32d
2 parents 3f5af2c + 6cfbe88
commit d8ba32d
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tensorflow/lite/kernels/split_v.cc b/tensorflow/lite/kernels/split_v.cc
@@ -96,6 +96,8 @@ TfLiteStatus ResizeOutputTensors(TfLiteContext* context, TfLiteNode* node,
     }
   }
 
+  TF_LITE_ENSURE(context, axis_value >= 0);
+  TF_LITE_ENSURE(context, axis_value < NumDimensions(input));
   const int input_size = SizeOfDimension(input, axis_value);
 
   if (minus_one_index != -1) {