From b49f212c1d1afa26f320c0bbdc136cecb0bd670f Mon Sep 17 00:00:00 2001 From: TristanInSec Date: Fri, 17 Apr 2026 15:27:56 -0400 Subject: [PATCH] Validate gather coordinate indices at runtime Replace TFLITE_DCHECK bounds checks with runtime validation that returns kTfLiteError when a coordinate index is negative or exceeds the axis size, consistent with the approach used in GatherNd. --- tensorflow/lite/micro/kernels/gather.cc | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tensorflow/lite/micro/kernels/gather.cc b/tensorflow/lite/micro/kernels/gather.cc index a0af4c0edda..003a741d7e6 100644 --- a/tensorflow/lite/micro/kernels/gather.cc +++ b/tensorflow/lite/micro/kernels/gather.cc @@ -82,13 +82,15 @@ TfLiteStatus Gather(const TfLiteGatherParams* params, for (int batch = 0; batch < batch_size; ++batch) { for (int outer = 0; outer < outer_size; ++outer) { for (int coord = 0; coord < coord_size; ++coord) { - TFLITE_DCHECK_GE(coords_data[coord], 0); - TFLITE_DCHECK_LT(coords_data[coord], axis_size); + CoordsT idx = coords_data[batch * coord_size + coord]; + if (idx < 0 || idx >= axis_size) { + return kTfLiteError; + } std::memcpy(output_data + (((batch * outer_size) + outer) * coord_size + coord) * inner_size, input_data + (((batch * outer_size) + outer) * axis_size + - coords_data[batch * coord_size + coord]) * + idx) * inner_size, sizeof(InputT) * inner_size); }