Merge pull request #49793 from geetachavan1/cherrypicks_AD5YQ

Prevent division by 0 in OneHot implementation
tensorflow · May 31, 2021 · e5d259f · e5d259f
2 parents a5fb5cb + e7848c9
commit e5d259f
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
@@ -67,6 +67,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;