Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix integer overflow in TFLite concat
PiperOrigin-RevId: 371013841
Change-Id: I6a4782ce7ca753e23ff31e7fb6aeb7f9d412cd29
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed Apr 28, 2021
1 parent 7b7352a commit 4253f96
Showing 1 changed file with 6 additions and 0 deletions.
6 changes: 6 additions & 0 deletions tensorflow/lite/kernels/concatenation.cc
Expand Up @@ -16,6 +16,8 @@ limitations under the License.

#include <stdint.h>

#include <limits>

#include "tensorflow/lite/c/builtin_op_data.h"
#include "tensorflow/lite/c/common.h"
#include "tensorflow/lite/kernels/internal/compatibility.h"
Expand Down Expand Up @@ -69,6 +71,10 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
TF_LITE_ENSURE_EQ(context, t->type, input_type);
for (int d = 0; d < t0->dims->size; ++d) {
if (d == axis) {
// Avoid integer overflow in sum_axis below
TF_LITE_ENSURE(context, t->dims->data[axis] >= 0);
TF_LITE_ENSURE(context, t->dims->data[axis] <=
std::numeric_limits<int>::max() - sum_axis);
sum_axis += t->dims->data[axis];
} else {
TF_LITE_ENSURE_EQ(context, t->dims->data[d], t0->dims->data[d]);
Expand Down

0 comments on commit 4253f96

Please sign in to comment.