[pull] master from tensorflow:master

tensorflow/compiler/mlir/tensorflow/BUILD

@@ -655,6 +655,7 @@ cc_library(
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/container:node_hash_map",
         "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:vlog_is_on",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:Analysis",
         "@llvm-project//mlir:FuncDialect",

tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.cc

@@ -27,6 +27,7 @@ limitations under the License.
 #include "absl/container/flat_hash_set.h"
 #include "absl/container/node_hash_map.h"
 #include "absl/log/log.h"
+#include "absl/log/vlog_is_on.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"

tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.cc

@@ -18,7 +18,6 @@ limitations under the License.
 #include <algorithm>
 #include <array>
 #include <cassert>
-#include <complex>
 #include <cstddef>
 #include <cstdint>
 #include <iterator>

tensorflow/core/distributed_runtime/tensor_coding.cc

@@ -167,6 +167,9 @@ bool TensorResponse::ParseTensorSubmessage(
                  .ok()) {
           return false;
         }
+        if (shape.num_elements() != 0) {
+          return false;
+        }
         Tensor t(allocator_, tensor_meta->dtype(), shape);
         tensor_ = std::move(t);
       }

tensorflow/core/distributed_runtime/tensor_coding_test.cc

@@ -180,6 +180,64 @@ TEST_F(TensorResponseTest, InitPartialOverflow) {
   EXPECT_TRUE(absl::IsInvalidArgument(s));
 }
 
+TEST_F(TensorResponseTest, ZeroLengthTensorMissingContentAccepted) {
+  RecvTensorResponse proto;
+  proto.set_is_dead(false);
+  proto.set_send_start_micros(123456);
+  TensorProto* tensor_proto = proto.mutable_tensor();
+  tensor_proto->set_dtype(DT_FLOAT);
+  tensor_proto->mutable_tensor_shape()->add_dim()->set_size(0);
+
+  std::string encoded;
+  proto.AppendToString(&encoded);
+  StringSource source(&encoded, 1024);
+  TensorResponse response;
+  DummyDevice cpu_device(Env::Default());
+  response.InitAlloc(&cpu_device, AllocatorAttributes());
+  EXPECT_TRUE(response.ParseFrom(&source).ok());
+  EXPECT_EQ(response.tensor().NumElements(), 0);
+}
+
+TEST_F(TensorResponseTest,
+       NonEmptyTensorMissingContentZeroInitializedViaSlowPath) {
+  RecvTensorResponse proto;
+  proto.set_is_dead(false);
+  proto.set_send_start_micros(123456);
+  TensorProto* tensor_proto = proto.mutable_tensor();
+  tensor_proto->set_dtype(DT_FLOAT);
+  tensor_proto->mutable_tensor_shape()->add_dim()->set_size(10);
+
+  std::string encoded;
+  proto.AppendToString(&encoded);
+  StringSource source(&encoded, 1024);
+  TensorResponse response;
+  DummyDevice cpu_device(Env::Default());
+  response.InitAlloc(&cpu_device, AllocatorAttributes());
+  EXPECT_TRUE(response.ParseFrom(&source).ok());
+  EXPECT_EQ(response.tensor().NumElements(), 10);
+  for (int i = 0; i < 10; ++i) {
+    EXPECT_EQ(response.tensor().flat<float>()(i), 0.0f);
+  }
+}
+
+TEST_F(TensorResponseTest, AmplificationProtoWithoutContentRejected) {
+  RecvTensorResponse proto;
+  proto.set_is_dead(false);
+  proto.set_send_start_micros(123456);
+  TensorProto* tensor_proto = proto.mutable_tensor();
+  tensor_proto->set_dtype(DT_FLOAT);
+  // 1 billion floats = 4GB (> 2GB safe limit) with no tensor_content
+  tensor_proto->mutable_tensor_shape()->add_dim()->set_size(1000000000);
+
+  std::string encoded;
+  proto.AppendToString(&encoded);
+  StringSource source(&encoded, 1024);
+  TensorResponse response;
+  DummyDevice cpu_device(Env::Default());
+  response.InitAlloc(&cpu_device, AllocatorAttributes());
+  EXPECT_FALSE(response.ParseFrom(&source).ok());
+}
+
 std::string MakeFloatTensorTestCase(int num_elems) {
   std::vector<int8_t> v(num_elems);
   for (int i = 0; i < num_elems; i++) {

tensorflow/core/kernels/data/cache_dataset_ops.cc

@@ -94,13 +94,13 @@ class DatasetRandomAccessCache {
                           GetIteratorResourceFromDataset(ctx, input_));
       TF_RETURN_IF_ERROR(iter_resource_->SetIteratorFromDataset(ctx, input_));
     }
-    if (index >= cache_.size()) {
-      TF_RETURN_IF_ERROR(ExtendTempCacheToIndex(index, ctx));
-    }
     if (index < 0) {
       return absl::InvalidArgumentError(
           absl::StrCat("Expected index >= 0; Received index: ", index));
     }
+    if (index >= static_cast<int64_t>(cache_.size())) {
+      TF_RETURN_IF_ERROR(ExtendTempCacheToIndex(index, ctx));
+    }
     *out_tensors = cache_.at(index);
     return absl::OkStatus();
   }
@@ -111,7 +111,7 @@ class DatasetRandomAccessCache {
  private:
   absl::Status ExtendTempCacheToIndex(int64_t index, OpKernelContext* ctx) {
     bool end_of_sequence;
-    while (cache_.size() <= index) {
+    while (static_cast<int64_t>(cache_.size()) <= index) {
       std::vector<Tensor> out_tensors;
       TF_RETURN_IF_ERROR(
           iter_resource_->GetNext(ctx, &out_tensors, &end_of_sequence));
@@ -151,15 +151,22 @@ class IteratorRandomAccessCache {
   explicit IteratorRandomAccessCache(const DatasetBase* input)
       : input_(input) {}
 
-  absl::Status Get(AnyContext ctx, size_t element_position,
+  absl::Status Get(AnyContext ctx, int64_t element_position,
                    std::vector<Tensor>* out_tensors) {
-    if (element_position < cache_.size() && !cache_[element_position].empty()) {
+    if (element_position < 0) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("Element position must be non-negative; Received: ",
+                       element_position));
+    }
+
+    if (element_position < static_cast<int64_t>(cache_.size()) &&
+        !cache_[element_position].empty()) {
       *out_tensors = cache_[element_position];
       return absl::OkStatus();
     }
 
     TF_RETURN_IF_ERROR(input_->Get(ctx, element_position, out_tensors));
-    if (element_position >= cache_.size()) {
+    if (element_position >= static_cast<int64_t>(cache_.size())) {
       cache_.resize(element_position + 1);
     }
     cache_[element_position] = *out_tensors;

tensorflow/lite/kernels/dynamic_update_slice.cc

@@ -278,15 +278,23 @@ void DynamicUpdateSliceInt4(const TfLiteTensor* input,
   // If the update is the entirety of the output, then simply copy it and
   // return.
   if (input_shape.FlatSize() == update_shape.FlatSize()) {
-    memcpy(output_data, update_data, input->bytes);
+    // Clamp the copy length to the destination size. `input->bytes` is derived
+    // from the raw FlatBuffer buffer size for constant tensors and may be
+    // inflated relative to the logical shape, so it must be bounded by
+    // `output->bytes` to avoid an out-of-bounds write.
+    memcpy(output_data, update_data, std::min(input->bytes, output->bytes));
     return;
   }
   RuntimeShape clamped_start_indices =
       ClampStartIndices(input_dims, indices_data, input_shape, update_shape);
 
   // If the operation is not done in-place, copy the input data to the output.
+  // Clamp the copy length to the destination size, since `input->bytes` may be
+  // larger than the output allocation for constant tensors loaded from a
+  // FlatBuffer model.
+  size_t bytes = std::min(input->bytes, output->bytes);
   if (input->data.data != output->data.data) {
-    memcpy(output->data.data, input->data.data, input->bytes);
+    memcpy(output->data.data, input->data.data, bytes);
   }
 
   // Update tensor has no elements. Skip.

tensorflow/lite/tools/cmake/modules/cpuinfo.cmake

@@ -41,4 +41,5 @@ set(CPUINFO_BUILD_BENCHMARKS OFF CACHE BOOL "Disable cpuinfo micro-benchmarks")
 add_subdirectory(
   "${cpuinfo_SOURCE_DIR}"
   "${cpuinfo_BINARY_DIR}"
+  EXCLUDE_FROM_ALL
 )

tensorflow/lite/tools/cmake/modules/eigen.cmake

@@ -99,4 +99,4 @@ set(EIGEN_TEST_SYCL OFF CACHE BOOL "Disable Sycl test")
 set(EIGEN_SYCL_TRISYCL OFF CACHE BOOL "Disable triSYCL test")
 # Make sure only MPL2.0 or more permissively licensed code is included.
 add_compile_definitions(EIGEN_MPL2_ONLY)
-add_subdirectory("${eigen_SOURCE_DIR}" "${eigen_BINARY_DIR}")
+add_subdirectory("${eigen_SOURCE_DIR}" "${eigen_BINARY_DIR}" EXCLUDE_FROM_ALL)

tensorflow/lite/tools/cmake/modules/farmhash.cmake

@@ -44,4 +44,5 @@ set(FARMHASH_SOURCE_DIR "${farmhash_SOURCE_DIR}" CACHE PATH
 add_subdirectory(
   "${CMAKE_CURRENT_LIST_DIR}/farmhash"
   "${farmhash_BINARY_DIR}"
+  EXCLUDE_FROM_ALL
 )

tensorflow/lite/tools/cmake/modules/fft2d.cmake

@@ -37,4 +37,5 @@ set(FFT2D_SOURCE_DIR "${fft2d_SOURCE_DIR}" CACHE PATH "fft2d source")
 add_subdirectory(
   "${CMAKE_CURRENT_LIST_DIR}/fft2d"
   "${fft2d_BINARY_DIR}"
+  EXCLUDE_FROM_ALL
 )

tensorflow/lite/tools/cmake/modules/flatbuffers.cmake

@@ -44,6 +44,7 @@ add_definitions(-DNOMINMAX=1)
 add_subdirectory(
   "${flatbuffers_SOURCE_DIR}"
   "${flatbuffers_BINARY_DIR}"
+  EXCLUDE_FROM_ALL
 )
 remove_definitions(-DNOMINMAX)
 

tensorflow/lite/tools/cmake/modules/gemmlowp.cmake

@@ -48,6 +48,7 @@ set(GEMMLOWP_SOURCE_DIR "${gemmlowp_SOURCE_DIR}" CACHE PATH "Source directory")
 add_subdirectory(
   "${gemmlowp_SOURCE_DIR}/contrib"
   "${gemmlowp_BINARY_DIR}"
+  EXCLUDE_FROM_ALL
 )
 
 set(BUILD_TESTING ${BUILD_TESTING_TMP})

tensorflow/lite/tools/cmake/modules/neon2sse.cmake

@@ -40,4 +40,5 @@ endif()
 add_subdirectory(
   "${neon2sse_SOURCE_DIR}"
   "${neon2sse_BINARY_DIR}"
+  EXCLUDE_FROM_ALL
 )

tensorflow/lite/tools/cmake/modules/ruy.cmake

@@ -37,4 +37,5 @@ set(RUY_SOURCE_DIR "${ruy_SOURCE_DIR}" CACHE PATH "RUY source directory")
 add_subdirectory(
   "${ruy_SOURCE_DIR}"
   "${ruy_BINARY_DIR}"
+  EXCLUDE_FROM_ALL
 )

tensorflow/python/data/kernel_tests/prefetch_test.py

@@ -25,6 +25,7 @@
 from tensorflow.python.data.ops import options as options_lib
 from tensorflow.python.data.ops import prefetch_op
 from tensorflow.python.framework import combinations
+from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors
 from tensorflow.python.ops import script_ops
@@ -40,6 +41,20 @@ def testBufferSize(self, buffer_size):
     dataset = dataset_ops.Dataset.range(10).prefetch(buffer_size=buffer_size)
     self.assertDatasetProduces(dataset, expected_output=range(10))
 
+  @combinations.generate(test_base.default_test_combinations())
+  def testTensorBufferSize(self):
+    # `buffer_size` is documented as an int64 scalar `tf.Tensor`.
+    dataset = dataset_ops.Dataset.range(10).prefetch(
+        buffer_size=constant_op.constant(2, dtypes.int64)
+    )
+    self.assertDatasetProduces(dataset, expected_output=range(10))
+
+    # A tensor holding the AUTOTUNE value keeps autotuning enabled.
+    dataset = dataset_ops.Dataset.range(10).prefetch(
+        buffer_size=constant_op.constant(-1, dtypes.int64)
+    )
+    self.assertDatasetProduces(dataset, expected_output=range(10))
+
   @combinations.generate(
       combinations.times(test_base.eager_only_combinations(),
                          combinations.combine(buffer_size=[0, 1, 2, 42])))

tensorflow/python/data/kernel_tests/shuffle_test.py

@@ -178,6 +178,20 @@ def testDefaultArguments(self):
     for i in range(5):
       self.assertEqual(10, counts[i])
 
+  @combinations.generate(test_base.default_test_combinations())
+  def testExcessiveBufferSize(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices([1, 2, 3, 4, 5])
+    with self.assertRaisesRegex(ValueError, "buffer_size"):
+      dataset.shuffle(buffer_size=sys.maxsize)
+    with self.assertRaisesRegex(ValueError, "buffer_size"):
+      dataset.shuffle(buffer_size=2**31 - 1)
+    with self.assertRaisesRegex(ValueError, "buffer_size"):
+      dataset.shuffle(buffer_size=np.int64(2**31 - 1))
+    with self.assertRaisesRegex(ValueError, "buffer_size"):
+      dataset.shuffle(
+          buffer_size=constant_op.constant(2**31 - 1, dtype=dtypes.int64)
+      )
+
   @combinations.generate(
       combinations.times(
           test_base.default_test_combinations(),

tensorflow/python/data/ops/dataset_ops.py

@@ -1062,6 +1062,12 @@ def zip(*args, datasets=None, name=None) -> "DatasetV2":
     >>> [(i.item(), j.item()) for i, j in ds.as_numpy_iterator()]
     [(1, 13), (2, 14)]
 
+    Note: `zip` iterates its input datasets in parallel. If one dataset has
+    been shuffled and another has not, the element correspondence from the
+    original order will be lost. To keep elements aligned after shuffling,
+    apply `shuffle` after `zip` rather than before it. See
+    `tf.data.Dataset.shuffle` for more details.
+
     Args:
       *args: Datasets or nested structures of datasets to zip together. This
         can't be set if `datasets` is set.
@@ -1492,6 +1498,21 @@ def shuffle(
     # [18, 4, 9, 2, 17, 8, 5, 10, 0, 6, 16, 3, 19, 7, 14, 11, 15, 13, 12, 1]
     ```
 
+    #### Using shuffle with zip
+
+    When you `shuffle` a dataset and then `zip` it with an unshuffled dataset,
+    the element pairing will not match the original order because each input
+    dataset is iterated independently. If you need to shuffle multiple datasets
+    while preserving the correspondence between their elements, you should
+    shuffle after zipping:
+
+    ```python
+    # Correct: shuffle after zipping to keep pairs aligned.
+    a = tf.data.Dataset.range(3)
+    b = tf.data.Dataset.range(3)
+    dataset = tf.data.Dataset.zip(a, b).shuffle(3)
+    ```
+
     Args:
       buffer_size: An int or `tf.int64` scalar `tf.Tensor`, representing the
         number of elements from this dataset from which the new dataset will

tensorflow/python/data/ops/prefetch_op.py

@@ -18,6 +18,7 @@
 from tensorflow.python.data.ops import debug_mode
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_util
 from tensorflow.python.ops import gen_dataset_ops
 
 
@@ -39,6 +40,14 @@ def __init__(self, input_dataset, buffer_size, slack_period=None, name=None):
     self._buffer_size = ops.convert_to_tensor(
         buffer_size, dtype=dtypes.int64, name="buffer_size")
     self._name = name
+    # `legacy_autotune` must be a Python bool, so decide from the statically
+    # known value of the converted `buffer_size`; it is None for symbolic
+    # tensors, which then use the non-legacy path.
+    buffer_size_constant = tensor_util.constant_value(self._buffer_size)
+    legacy_autotune = (
+        buffer_size_constant is not None
+        and int(buffer_size_constant) == dataset_ops.AUTOTUNE
+    )
     # pylint: disable=protected-access
     # We colocate the prefetch dataset with its input as this collocation only
     # happens automatically in graph mode.
@@ -47,6 +56,7 @@ def __init__(self, input_dataset, buffer_size, slack_period=None, name=None):
           input_dataset._variant_tensor,
           buffer_size=self._buffer_size,
           slack_period=slack_period,
-          legacy_autotune=(buffer_size == dataset_ops.AUTOTUNE),
-          **self._common_args)
+          legacy_autotune=legacy_autotune,
+          **self._common_args
+      )
     super().__init__(input_dataset, variant_tensor)

tensorflow/python/data/ops/shuffle_op.py

@@ -19,8 +19,16 @@
 from tensorflow.python.eager import context
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_util
 from tensorflow.python.ops import gen_dataset_ops
 
+# Sanity limit on the number of elements in the shuffle buffer. The C++
+# kernel eagerly allocates a slot for every element up to `buffer_size` when
+# the iterator is created, so a pathologically large value (e.g.
+# `sys.maxsize`) reaches that allocation and crashes the process instead of
+# raising a catchable error.
+_MAX_SHUFFLE_BUFFER_SIZE_ELEMENTS = 1 << 30  # ~1 billion elements
+
 
 def _shuffle(  # pylint: disable=unused-private-name
     input_dataset,
@@ -48,6 +56,18 @@ def __init__(
     self._input_dataset = input_dataset
     self._buffer_size = ops.convert_to_tensor(
         buffer_size, dtype=dtypes.int64, name="buffer_size")
+    constant_buffer_size = tensor_util.constant_value(self._buffer_size)
+    if (
+        constant_buffer_size is not None
+        and constant_buffer_size > _MAX_SHUFFLE_BUFFER_SIZE_ELEMENTS
+    ):
+      raise ValueError(
+          "`buffer_size` must not exceed "
+          f"{_MAX_SHUFFLE_BUFFER_SIZE_ELEMENTS} elements, but got "
+          f"{constant_buffer_size}. Requesting a shuffle buffer this large "
+          "would cause the dataset to abort the process instead of raising "
+          "a catchable error."
+      )
     self._seed, self._seed2 = random_seed.get_seed(seed)
     self._reshuffle_each_iteration = reshuffle_each_iteration
     self._name = name