Fix out of bound error in ReverseSequence Op shape function

PiperOrigin-RevId: 411896080
Change-Id: I7e59a38e2f960886edf2b6c54ed5a84e86a9b193

tensorflow/core/ops/array_ops.cc

@@ -1653,11 +1653,21 @@ REGISTER_OP("ReverseSequence")
         return errors::InvalidArgument(
             "batch_dim must be < input rank: ", batch_dim, " vs. ", input_rank);
       }
+
       if (seq_dim >= input_rank) {
         return errors::InvalidArgument(
             "seq_dim must be < input rank: ", seq_dim, " vs. ", input_rank);
       }
 
+      // To prevent out of bound access when calling c->Dim(input, batch_dim),
+      // batch_dim range [-1 * input rank, input rank) is allowed. However,
+      // the op implementation has a stricter bound for batch_dim requiring >= 0
+      // value. Thus, perform strict check here.
+      if (batch_dim < 0) {
+        return errors::InvalidArgument("batch_dim must be >=0, got ",
+                                       batch_dim);
+      }
+
       DimensionHandle batch_dim_dim = c->Dim(input, batch_dim);
       TF_RETURN_IF_ERROR(
           c->Merge(batch_dim_dim, c->Dim(seq_lens_shape, 0), &batch_dim_dim));