Merge pull request #49548 from geetachavan1/cherrypicks_SBBYG

Fix FPE issue with `tf.raw_ops.Reverse`.
tensorflow · May 27, 2021 · 5cf9b08 · 5cf9b08
2 parents a525ac6 + 9b1956d
commit 5cf9b08
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/tensorflow/core/kernels/reverse_op.cc b/tensorflow/core/kernels/reverse_op.cc
@@ -158,6 +158,12 @@ class ReverseOp : public OpKernel {
 
   void Compute(OpKernelContext* context) override {
     const Tensor& input = context->input(0);
+    // If input is provided, check to make sure the first dimension is valid.
+    if (input.dims() > 0) {
+      OP_REQUIRES(
+          context, input.dim_size(0) != 0,
+          errors::InvalidArgument("Invalid input first dimension. Found 0."));
+    }
     const Tensor& dims = context->input(1);
 
     if (TensorShapeUtils::IsScalar(input.shape())) {