Merge pull request #126158 from vishwakt:fix-slice-begin-size-overflow

PiperOrigin-RevId: 1005385415

tensorflow/core/kernels/slice_op.cc

@@ -105,7 +105,12 @@ void SharedSliceValidation(OpKernelContext* context, const Tensor& input,
                   absl::InvalidArgumentError(
                       absl::StrCat("Expected begin[", i, "] in [0, ",
                                    input.dim_size(i), "], but got ", b)));
-      OP_REQUIRES(context, 0 <= s && b + s <= input.dim_size(i),
+      // `b` has been validated to lie in [0, dim_size(i)] just above, so
+      // `dim_size(i) - b` cannot underflow. Comparing against it instead of
+      // computing `b + s` avoids a signed overflow: for a large `dim_size(i)`,
+      // `b + s` can wrap negative and pass the upper-bound check, letting the
+      // wrapped limit reach a fatal `Tensor::Slice` invariant downstream.
+      OP_REQUIRES(context, 0 <= s && s <= input.dim_size(i) - b,
                   absl::InvalidArgumentError(absl::StrCat(
                       "Expected size[", i, "] in [0, ", input.dim_size(i) - b,
                       "], but ", "got ", s)));

tensorflow/python/kernel_tests/array_ops/BUILD

@@ -714,6 +714,7 @@ cuda_py_strict_test(
         "//tensorflow/python/framework:for_generated_wrappers",
         "//tensorflow/python/framework:test_lib",
         "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:array_ops_gen",
         "//tensorflow/python/ops:gradients_impl",
         "//tensorflow/python/ops:math_ops",
         "//tensorflow/python/ops:nn_ops",

tensorflow/python/kernel_tests/array_ops/slice_op_test.py

@@ -24,6 +24,7 @@ from tensorflow.python.framework import errors_impl
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_array_ops
 from tensorflow.python.ops import gradients_impl
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import nn_ops
@@ -211,6 +212,32 @@ class SliceTest(test.TestCase):
         TypeError, "must be integers or slices, not tuple"):
       self.evaluate(func(input_val))
 
+  @test_util.run_in_graph_and_eager_modes
+  @test_util.disable_xla(
+      "XLA shape inference rejects the reshape to an INT64_MAX dimension, so "
+      "the test cannot reach the Slice kernel under XLA"
+  )
+  def testBeginSizeOverflowRaises(self):
+    # Regression test for GitHub issue 126125. `begin + size` was validated
+    # with signed int64 addition, so for a large first dimension the sum could
+    # wrap negative, pass the upper-bound check, and reach a fatal
+    # `Tensor::Slice` invariant. It must raise instead.
+    i64_max = (1 << 63) - 1
+    value = array_ops.reshape(
+        constant_op.constant([], dtype=dtypes.float32),
+        constant_op.constant([i64_max, 0], dtype=dtypes.int64),
+    )
+    with self.assertRaisesRegex(
+        errors_impl.InvalidArgumentError, "Expected size"
+    ):
+      self.evaluate(
+          gen_array_ops.Slice(
+              input=value,
+              begin=constant_op.constant([i64_max, 0], dtype=dtypes.int64),
+              size=constant_op.constant([i64_max, 0], dtype=dtypes.int64),
+          )
+      )
+
   def testOutOfBoundsSizeWithConstantInput(self):
     # An out-of-bounds size is rejected by the Slice kernel in eager execution,
     # and must be rejected the same way inside a tf.function. Grappler folded a