Prevent division by 0 in TFLite

PiperOrigin-RevId: 370800311 Change-Id: I21ccdbd31c30118acc67df8751807ee2e0b12f91
tensorflow · Apr 28, 2021 · 106d8f4 · 106d8f4
1 parent 953f28d
commit 106d8f4
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 0 deletions.
diff --git a/tensorflow/lite/kernels/depth_to_space.cc b/tensorflow/lite/kernels/depth_to_space.cc
@@ -61,6 +61,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
 
   const int block_size = params->block_size;
+  TF_LITE_ENSURE(context, block_size > 0);
   const int input_height = input->dims->data[1];
   const int input_width = input->dims->data[2];
   const int input_channels = input->dims->data[3];

diff --git a/tensorflow/lite/kernels/depth_to_space_test.cc b/tensorflow/lite/kernels/depth_to_space_test.cc
@@ -60,6 +60,11 @@ TEST(DepthToSpaceOpModel, BadBlockSize) {
   EXPECT_DEATH(DepthToSpaceOpModel({TensorType_FLOAT32, {1, 1, 1, 4}}, 4),
                "Cannot allocate tensors");
 }
+
+TEST(DepthToSpaceOpModel, NoBlockSize) {
+  EXPECT_DEATH(DepthToSpaceOpModel({TensorType_FLOAT32, {1, 1, 1, 4}}, 0),
+               "Cannot allocate tensors");
+}
 #endif
 
 TEST(DepthToSpaceOpModel, Float32) {

diff --git a/tensorflow/lite/micro/kernels/depth_to_space.cc b/tensorflow/lite/micro/kernels/depth_to_space.cc
@@ -54,6 +54,7 @@ TfLiteStatus CalculateOpData(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
 
   const int block_size = params->block_size;
+  TF_LITE_ENSURE(context, block_size > 0);
   const int input_height = input->dims->data[kHeightRank];
   const int input_width = input->dims->data[kWidthRank];
   const int input_channels = input->dims->data[kDepthRank];