[pull] master from tensorflow:master

ci/official/utilities/extract_resultstore_links.py

@@ -114,8 +114,7 @@ def parse_log(file_path: str,
       else:
         tests_failed = re.search(TESTS_FAILED_RE, backtrack_line)
       if build_failed or tests_failed:
-        log_fragment = '\n'.join(
-            log_lines[max(k - 20, 0):min(end_line + 1, len(log_lines) - 1)])
+        log_fragment = '\n'.join(log_lines[max(k - 20, 0) : end_line + 1])
         lines['log_fragment'] = log_fragment
         lines['status'] = (InvokeStatus.build_failed if build_failed
                            else InvokeStatus.tests_failed)

tensorflow/core/kernels/conv_grad_filter_ops_3d.cc

@@ -17,6 +17,7 @@ limitations under the License.
 #define EIGEN_USE_THREADS
 
 #include <algorithm>
+#include <limits>
 #include <string>
 #include <utility>
 #include <vector>
@@ -880,6 +881,14 @@ void LaunchConvBackpropFilterOpImpl(
           : TensorShape({filter_shape.dim_size(4), dims.filter_size(0),
                          dims.filter_size(1), dims.filter_size(2),
                          filter_shape.dim_size(3)});
+  // Validate filter element count before allocation to prevent OOM on invalid
+  // inputs. GPU transformation uses 32-bit indexing via To32Bit().
+  OP_REQUIRES(
+      context,
+      filter_backprop->NumElements() <= std::numeric_limits<int32>::max(),
+      errors::InvalidArgument("Filter tensor num elements (",
+                              filter_backprop->NumElements(),
+                              ") exceeds 32-bit limit for GPU transformation"));
   OP_REQUIRES_OK(context,
                  context->allocate_temp(DataTypeToEnum<T>::value, dst_shape,
                                         &pre_transformed_filter_backprop));

tensorflow/core/kernels/conv_grad_filter_ops_launcher.cc

@@ -17,6 +17,7 @@ limitations under the License.
 #define EIGEN_USE_THREADS
 
 #include <algorithm>
+#include <limits>
 #include <utility>
 #include <vector>
 
@@ -399,6 +400,14 @@ void LaunchConv2DBackpropFilterOpImpl(
   // We compute filter backprop into temporary tensor, and then convert it to
   // the HWIO data format at the end.
 
+  // Validate filter element count before allocation to prevent OOM on invalid
+  // inputs. GPU transformation uses 32-bit indexing via To32Bit().
+  OP_REQUIRES(
+      ctx, filter_backprop->NumElements() <= std::numeric_limits<int32>::max(),
+      errors::InvalidArgument("Filter tensor num elements (",
+                              filter_backprop->NumElements(),
+                              ") exceeds 32-bit limit for GPU transformation"));
+
   Tensor pre_transformed_filter_backprop;
   OP_REQUIRES_OK(
       ctx,

tensorflow/core/kernels/conv_grad_input_ops.cc

@@ -17,6 +17,7 @@ limitations under the License.
 
 #include "tensorflow/core/kernels/conv_grad_input_ops.h"
 
+#include <limits>
 #include <utility>
 
 #include "tensorflow/core/profiler/lib/scoped_annotation.h"
@@ -291,6 +292,11 @@ void LaunchConv2DBackpropInputOpGpuImpl(
             : TensorShape({filter.dim_size(3), filter.dim_size(0),
                            filter.dim_size(1), filter.dim_size(2)});
 
+    if (filter.NumElements() > std::numeric_limits<int32>::max()) {
+      return errors::InvalidArgument(
+          "Filter tensor num elements (", filter.NumElements(),
+          ") exceeds 32-bit limit for GPU transformation");
+    }
     TF_RETURN_IF_ERROR(ctx->allocate_temp(DataTypeToEnum<T>::value, dst_shape,
                                           &transformed_filter));
     functor::TransformFilter<GPUDevice, T, int, 4>()(

tensorflow/core/kernels/conv_grad_input_ops_3d.cc

@@ -17,6 +17,7 @@ limitations under the License.
 #define EIGEN_USE_THREADS
 
 #include <algorithm>
+#include <limits>
 #include <string>
 #include <utility>
 #include <vector>
@@ -873,6 +874,12 @@ void LaunchConvBackpropInputOpImpl(
           : TensorShape({filter_shape.dim_size(4), dims.filter_size(0),
                          dims.filter_size(1), dims.filter_size(2),
                          filter_shape.dim_size(3)});
+  OP_REQUIRES(context,
+              filter.NumElements() <= std::numeric_limits<int32>::max(),
+              errors::InvalidArgument(
+                  "Filter tensor num elements (", filter.NumElements(),
+                  ") exceeds 32-bit limit for GPU transformation"));
+
   OP_REQUIRES_OK(context,
                  context->allocate_temp(DataTypeToEnum<T>::value, dst_shape,
                                         &transformed_filter));

tensorflow/core/kernels/conv_ops_fused_impl.h

@@ -38,6 +38,7 @@ limitations under the License.
 #define EIGEN_USE_GPU
 #endif  // GOOGLE_CUDA
 
+#include <limits>
 #include <string>
 #include <type_traits>
 #include <utility>
@@ -558,8 +559,14 @@ struct LaunchFusedConv2DOp<GPUDevice, T> {
               : TensorShape({filter.dim_size(3), filter.dim_size(0),
                              filter.dim_size(1), filter.dim_size(2)});
 
+      if (filter.NumElements() > std::numeric_limits<int32>::max()) {
+        return errors::InvalidArgument(
+            "Filter tensor num elements (", filter.NumElements(),
+            ") exceeds 32-bit limit for GPU transformation");
+      }
       TF_RETURN_IF_ERROR(context->allocate_temp(
           DataTypeToEnum<T>::value, dst_shape, &transformed_filter));
+
       functor::TransformFilter<GPUDevice, T, int, 4>()(
           context->eigen_device<GPUDevice>(), dst_format,
           To32Bit(filter.tensor<T, 4>()),

tensorflow/core/kernels/conv_ops_impl.h

@@ -1133,6 +1133,12 @@ void LaunchConvOpImpl(OpKernelContext* context, bool cudnn_use_autotune,
     }
   }
   TensorShape dst_shape(dst_shape_vec);
+  OP_REQUIRES(context,
+              filter.NumElements() <= std::numeric_limits<int32>::max(),
+              errors::InvalidArgument(
+                  "Filter tensor num elements (", filter.NumElements(),
+                  ") exceeds 32-bit limit for GPU transformation"));
+
   OP_REQUIRES_OK(context,
                  context->allocate_temp(DataTypeToEnum<T>::value, dst_shape,
                                         &transformed_filter));

tensorflow/core/kernels/deserialize_sparse_string_op.cc

@@ -203,7 +203,7 @@ class DeserializeSparseOp : public OpKernel {
       target_shape.vec<int64_t>()(i) = serialized_sparse.shape().dim_size(i);
     }
     for (int i = 0; i < output.dims() - 1; ++i) {
-      target_shape.vec<int64_t>()(i + ndims - 1) = output.shape().data()[i + 1];
+      target_shape.vec<int64_t>()(i + ndims - 1) = output.shape()[i + 1];
     }
 
     ReshapeSparseTensor<CPUDevice>(context, output.indices(), input_shape,

tensorflow/core/kernels/image/crop_and_resize_op.cc

@@ -55,24 +55,21 @@ using Callback = std::function<void()>;
 static inline absl::Status ParseAndCheckBoxSizes(const Tensor& boxes,
                                                  const Tensor& box_index,
                                                  int* num_boxes) {
-  if (boxes.NumElements() == 0 && box_index.NumElements() == 0) {
-    *num_boxes = 0;
-    return absl::OkStatus();
-  }
-  // The shape of 'boxes' is [num_boxes, 4].
+  // The shape of 'boxes' is [num_boxes, 4] and the shape of 'box_index' is
+  // [num_boxes]. The ranks must be validated even when both tensors are
+  // empty, since the kernels later access them as rank-2 and rank-1 tensors.
   if (boxes.dims() != 2) {
     return absl::InvalidArgumentError(
-        absl::StrCat("boxes must be 2-D", boxes.shape().DebugString()));
+        absl::StrCat("boxes must be 2-D, got ", boxes.shape().DebugString()));
+  }
+  if (box_index.dims() != 1) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "box_index must be 1-D, got ", box_index.shape().DebugString()));
   }
   *num_boxes = boxes.dim_size(0);
   if (boxes.dim_size(1) != 4) {
     return absl::InvalidArgumentError("boxes must have 4 columns");
   }
-  // The shape of 'box_index' is [num_boxes].
-  if (box_index.dims() != 1) {
-    return absl::InvalidArgumentError(
-        absl::StrCat("box_index must be 1-D", box_index.shape().DebugString()));
-  }
   if (box_index.dim_size(0) != *num_boxes) {
     return absl::InvalidArgumentError("box_index has incompatible shape");
   }

tensorflow/core/kernels/split_v_op.cc

@@ -27,6 +27,7 @@ limitations under the License.
 #define PLUGGABLE_DEVICE_SUPPORTED_MACOS 1
 #endif
 
+#include <limits>
 #include <numeric>
 
 #include "unsupported/Eigen/CXX11/Tensor"  // from @eigen_archive
@@ -100,7 +101,15 @@ class SplitVOpBase : public OpKernel {
             "-input rank(-", input.dims(), ") <= split_dim < input rank (",
             input.dims(), "), but got ", split_dim_orig)));
 
-    Tlen input_size_split_dim = input_shape.dim_size(split_dim);
+    // Check that the input size fits in Tlen before converting it. Otherwise
+    // an int32 or int8 Tlen truncates it, and split sizes that sum to the
+    // truncated size silently drop the rest of the input.
+    const int64_t actual_input_size = input_shape.dim_size(split_dim);
+    OP_REQUIRES(context, actual_input_size <= std::numeric_limits<Tlen>::max(),
+                errors::InvalidArgument(
+                    "Input size along split_dim must be <= max(Tlen). Got: ",
+                    actual_input_size));
+    Tlen input_size_split_dim = static_cast<Tlen>(actual_input_size);
 
     // Special case 1: num_split == 1. Nothing to do.
     if (num_split == 1) {
@@ -128,6 +137,24 @@ class SplitVOpBase : public OpKernel {
                                        "input."));
         neg_one_dim = d;
       } else {
+        // Reject a negative size before summing it, so that
+        // 0 <= determined_size <= input_size_split_dim below. Otherwise a
+        // large negative size, such as the minimum of Tlen next to a -1, makes
+        // `input_size_split_dim - determined_size` overflow.
+        OP_REQUIRES(context, size >= 0,
+                    errors::InvalidArgument("Split size at index ", d,
+                                            " must be >= 0. Got: ", size));
+        // Accumulate with an explicit overflow guard. `determined_size += size`
+        // wraps for large `size_splits`, and a wrapped total can equal
+        // `input_size_split_dim` and pass the check below, letting the aligned
+        // slicing path compute an endpoint that reaches a fatal `Tensor::Slice`
+        // invariant. Rejecting overflow here also keeps that later path safe,
+        // since the accepted total then bounds every partial sum.
+        OP_REQUIRES(context,
+                    determined_size <= std::numeric_limits<Tlen>::max() - size,
+                    errors::InvalidArgument(
+                        "Sum of size_splits overflows the index type at index ",
+                        d, "."));
         determined_size += size;
       }
     }
@@ -147,13 +174,6 @@ class SplitVOpBase : public OpKernel {
       (*split_sizes_vec)[neg_one_dim] = input_size_split_dim - determined_size;
     }
 
-    for (int i = 0; i < split_sizes_vec->size(); ++i) {
-      const Tlen& split_size = (*split_sizes_vec)[i];
-      OP_REQUIRES(context, split_size >= Tlen(0),
-                  errors::InvalidArgument("Split size at index ", i,
-                                          " must be >= 0. Got: ", split_size));
-    }
-
     // Special case 2: split along the 1st dimension. The requirements are that
     // either we are splitting the outer dimension of two or more such that
     // every outer subpart is aligned or that the split sizes mean that they are

tensorflow/core/lib/db/sqlite_test.cc

@@ -170,7 +170,7 @@ TEST_F(SqliteTest, UnsafeColumn) {
   stmt = db_->PrepareOrDie("SELECT b FROM T ORDER BY a");
   TF_ASSERT_OK(stmt.Step(&is_done_));
   absl::string_view p = stmt.ColumnStringUnsafe(0);
-  EXPECT_EQ('h', *p.data());
+  EXPECT_EQ('h', p[0]);
   TF_ASSERT_OK(stmt.Step(&is_done_));
   // This will actually happen, but it's not safe to test this behavior.
   // EXPECT_EQ('t', *p.data());

tensorflow/python/kernel_tests/array_ops/split_op_test.py

@@ -120,6 +120,63 @@ def testExplicitNum(self):
     self.assertAllEqual(r[1], value[2:4])
     self.assertAllEqual(r[2], value[4:])
 
+  @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 SplitV kernel under XLA"
+  )
+  def testSizeSplitsOverflowRaises(self):
+    # Regression test for GitHub issue 126126. The cumulative sum of
+    # size_splits was computed with unchecked signed addition, so a wrapped
+    # total could equal the input dimension, pass validation, and reach a
+    # fatal `Tensor::Slice` invariant in the aligned slicing path. It must
+    # raise instead.
+    i64_max = (1 << 63) - 1
+    i32_max = (1 << 31) - 1
+    for input_size, size_splits, dtype, message in (
+        (i64_max, [i64_max, i64_max, i64_max, 2], dtypes.int64, "overflow"),
+        # A -1 does not keep the other sizes from overflowing.
+        (i64_max, [-1, i64_max, i64_max, 2], dtypes.int64, "overflow"),
+        # int32 sizes overflow at their own width in the kernel. In graph
+        # mode the shape function, which sums in int64, rejects the mismatch
+        # before the kernel runs.
+        (5, [i32_max, i32_max, 5], dtypes.int32, "overflow|can't split axis"),
+    ):
+      with self.subTest(size_splits=size_splits, dtype=dtype.name):
+        value = array_ops.reshape(
+            constant_op.constant([], dtype=dtypes.float32),
+            constant_op.constant([input_size, 0], dtype=dtypes.int64),
+        )
+        with self.assertRaisesRegex(
+            (ValueError, errors_impl.InvalidArgumentError), message
+        ):
+          self.evaluate(
+              array_ops.split(
+                  value, constant_op.constant(size_splits, dtype=dtype), axis=0
+              )
+          )
+
+  @test_util.run_in_graph_and_eager_modes
+  @test_util.disable_xla("Checks the SplitV kernel, which XLA replaces")
+  def testInputSizeAboveSizeSplitsTypeRaises(self):
+    # With int32 size_splits, an input size above INT32_MAX was truncated to
+    # int32, so [1, 2] matched an input of size 2**32 + 3 and the split
+    # silently dropped the rest of the input. In graph mode the shape
+    # function, which compares in int64, rejects the mismatch first.
+    value = array_ops.reshape(
+        constant_op.constant([], dtype=dtypes.float32),
+        constant_op.constant([(1 << 32) + 3, 0], dtype=dtypes.int64),
+    )
+    with self.assertRaisesRegex(
+        (ValueError, errors_impl.InvalidArgumentError),
+        r"must be <= max\(Tlen\)|can't split axis",
+    ):
+      self.evaluate(
+          array_ops.split(
+              value, constant_op.constant([1, 2], dtype=dtypes.int32), axis=0
+          )
+      )
+
   @test_util.run_in_graph_and_eager_modes
   def testListOfScalarTensors(self):
     a = math_ops.cast(5, dtypes.int32)