From 417494c5ee96dc2e609a0767f9c775698bc232b1 Mon Sep 17 00:00:00 2001 From: Vishwak Thatikonda Date: Tue, 25 Aug 2026 23:57:38 -0700 Subject: [PATCH 1/3] Reject Slice begin+size overflow instead of aborting --- tensorflow/core/kernels/slice_op.cc | 7 ++++++- .../python/kernel_tests/array_ops/BUILD | 1 + .../kernel_tests/array_ops/slice_op_test.py | 19 +++++++++++++++++++ 3 files changed, 26 insertions(+), 1 deletion(-) diff --git a/tensorflow/core/kernels/slice_op.cc b/tensorflow/core/kernels/slice_op.cc index 5cef063c379..ed5b9ceb7f0 100644 --- a/tensorflow/core/kernels/slice_op.cc +++ b/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))); diff --git a/tensorflow/python/kernel_tests/array_ops/BUILD b/tensorflow/python/kernel_tests/array_ops/BUILD index 72f6e96c5a8..692c6ee11a3 100644 --- a/tensorflow/python/kernel_tests/array_ops/BUILD +++ b/tensorflow/python/kernel_tests/array_ops/BUILD @@ -702,6 +702,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", diff --git a/tensorflow/python/kernel_tests/array_ops/slice_op_test.py b/tensorflow/python/kernel_tests/array_ops/slice_op_test.py index 95a20b70bca..f4bac312267 100644 --- a/tensorflow/python/kernel_tests/array_ops/slice_op_test.py +++ b/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,24 @@ 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 + 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 _testSliceMatrixDim0(self, x, begin, size): tf_ans = self.evaluate(array_ops.slice(x, [begin, 0], [size, x.shape[1]])) np_ans = x[begin:begin + size, :] From 883f691c57c60a8d8b47052aee75ae72d9370e94 Mon Sep 17 00:00:00 2001 From: Vishwak Thatikonda Date: Wed, 26 Aug 2026 09:35:10 -0700 Subject: [PATCH 2/3] Skip the Slice overflow test when XLA is enabled --- tensorflow/python/kernel_tests/array_ops/slice_op_test.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tensorflow/python/kernel_tests/array_ops/slice_op_test.py b/tensorflow/python/kernel_tests/array_ops/slice_op_test.py index f4bac312267..4d9d5411e18 100644 --- a/tensorflow/python/kernel_tests/array_ops/slice_op_test.py +++ b/tensorflow/python/kernel_tests/array_ops/slice_op_test.py @@ -218,6 +218,10 @@ class SliceTest(test.TestCase): # 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. + if test_util.is_xla_enabled(): + # XLA cannot compile the reshape to an INT64_MAX dimension, so the test + # cannot reach the Slice kernel under XLA. + self.skipTest("XLA does not support shapes whose element count overflows") i64_max = (1 << 63) - 1 value = array_ops.reshape( constant_op.constant([], dtype=dtypes.float32), From b1a7b57a4572e15cbbef862363ee30785ecdb09d Mon Sep 17 00:00:00 2001 From: Vishwak Thatikonda Date: Wed, 26 Aug 2026 09:38:04 -0700 Subject: [PATCH 3/3] Use disable_xla decorator for the Slice overflow test --- tensorflow/python/kernel_tests/array_ops/slice_op_test.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tensorflow/python/kernel_tests/array_ops/slice_op_test.py b/tensorflow/python/kernel_tests/array_ops/slice_op_test.py index 4d9d5411e18..141d8be2373 100644 --- a/tensorflow/python/kernel_tests/array_ops/slice_op_test.py +++ b/tensorflow/python/kernel_tests/array_ops/slice_op_test.py @@ -213,15 +213,14 @@ class SliceTest(test.TestCase): 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. - if test_util.is_xla_enabled(): - # XLA cannot compile the reshape to an INT64_MAX dimension, so the test - # cannot reach the Slice kernel under XLA. - self.skipTest("XLA does not support shapes whose element count overflows") i64_max = (1 << 63) - 1 value = array_ops.reshape( constant_op.constant([], dtype=dtypes.float32),