Merge pull request #49086 from geetachavan1/cherrypicks_VQH2G

Prevent division by 0 in OneHot implementation
tensorflow · May 11, 2021 · b8679da · b8679da
2 parents c13f5f5 + 4523630
commit b8679da
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/tensorflow/lite/kernels/one_hot.cc b/tensorflow/lite/kernels/one_hot.cc
@@ -69,6 +69,11 @@ void OneHotComputeImpl(const OneHotContext& op_context) {
   for (int i = 0; i < op_context.axis; ++i) {
     prefix_dim_size *= op_context.indices->dims->data[i];
   }
+  if (prefix_dim_size == 0) {
+    // If indices tensor is degenerate, return a degenerate tensor, just like
+    // TensorFlow does.
+    return;
+  }
   const int suffix_dim_size = NumElements(op_context.indices) / prefix_dim_size;
   const int depth = *op_context.depth->data.i32;