[pull] master from tensorflow:master

tensorflow/compiler/tests/ternary_ops_test.py

@@ -60,18 +60,21 @@ def testLinspace(self, start, end, num):
     self.assertEqual(result[0], expected[0])
 
   def testRange(self):
-    self._testTernary(
-        math_ops.range,
-        np.int32(1),
-        np.int32(2),
-        np.int32(1),
-        expected=np.array([1], dtype=np.int32))
-    self._testTernary(
-        math_ops.range,
-        np.int32(1),
-        np.int32(7),
-        np.int32(2),
-        expected=np.array([1, 3, 5], dtype=np.int32))
+    for dtype in (self.int_types | self.float_types) - {np.uint8}:
+      self._testTernary(
+          math_ops.range,
+          dtype(1),
+          dtype(2),
+          dtype(1),
+          expected=np.array([1], dtype=dtype),
+      )
+      self._testTernary(
+          math_ops.range,
+          dtype(1),
+          dtype(7),
+          dtype(2),
+          expected=np.array([1, 3, 5], dtype=dtype),
+      )
 
   def testSelect(self):
     for dtype in self.numeric_types:

tensorflow/compiler/tf2xla/kernels/sequence_ops.cc

@@ -48,10 +48,10 @@ absl::StatusOr<xla::XlaOp> CreateRangeTensor(
   T limit = limit_literal.Get<T>({});
   T delta = delta_literal.Get<T>({});
 
-  if (delta == 0) {
+  if (delta == static_cast<T>(0)) {
     return errors::InvalidArgument("Requires delta != 0: ", delta);
   }
-  if (delta > 0) {
+  if (delta > static_cast<T>(0)) {
     if (start > limit) {
       return errors::InvalidArgument(
           "Requires start <= limit when delta > 0: ", start, "/", limit);
@@ -62,13 +62,21 @@ absl::StatusOr<xla::XlaOp> CreateRangeTensor(
           "Requires start >= limit when delta < 0: ", start, "/", limit);
     }
   }
-  int64_t size =
-      (std::is_integral<T>::value
-           ? static_cast<T>(
-                 limit == start
-                     ? 0
-                     : (std::abs(limit - start) - 1) / std::abs(delta) + 1)
-           : std::ceil(std::abs((limit - start) / delta)));
+  int64_t size;
+  if constexpr (std::is_integral<T>::value) {
+    int64_t start_i = static_cast<int64_t>(start);
+    int64_t limit_i = static_cast<int64_t>(limit);
+    int64_t delta_i = static_cast<int64_t>(delta);
+    size = (limit_i == start_i
+                ? 0
+                : (std::abs(limit_i - start_i) - 1) / std::abs(delta_i) + 1);
+  } else {
+    double start_f = static_cast<double>(start);
+    double limit_f = static_cast<double>(limit);
+    double delta_f = static_cast<double>(delta);
+    size = static_cast<int64_t>(
+        std::ceil(std::abs((limit_f - start_f) / delta_f)));
+  }
 
   return xla::ConstantR0(builder, start) +
          xla::ConstantR0(builder, delta) *
@@ -103,6 +111,13 @@ class RangeOp : public XlaOpKernel {
     DataType type = input_type(0);
     absl::StatusOr<xla::XlaOp> output;
     switch (type) {
+      case DT_INT8:
+        output = CreateRangeTensor<int8_t>(start, limit, delta, ctx->builder());
+        break;
+      case DT_INT16:
+        output =
+            CreateRangeTensor<int16_t>(start, limit, delta, ctx->builder());
+        break;
       case DT_INT32:
         output =
             CreateRangeTensor<int32_t>(start, limit, delta, ctx->builder());
@@ -111,6 +126,26 @@ class RangeOp : public XlaOpKernel {
         output =
             CreateRangeTensor<int64_t>(start, limit, delta, ctx->builder());
         break;
+      case DT_UINT16:
+        output =
+            CreateRangeTensor<uint16_t>(start, limit, delta, ctx->builder());
+        break;
+      case DT_UINT32:
+        output =
+            CreateRangeTensor<uint32_t>(start, limit, delta, ctx->builder());
+        break;
+      case DT_UINT64:
+        output =
+            CreateRangeTensor<uint64_t>(start, limit, delta, ctx->builder());
+        break;
+      case DT_HALF:
+        output =
+            CreateRangeTensor<Eigen::half>(start, limit, delta, ctx->builder());
+        break;
+      case DT_BFLOAT16:
+        output =
+            CreateRangeTensor<bfloat16>(start, limit, delta, ctx->builder());
+        break;
       case DT_FLOAT:
         output = CreateRangeTensor<float>(start, limit, delta, ctx->builder());
         break;
@@ -133,7 +168,7 @@ class RangeOp : public XlaOpKernel {
       xla::XlaOp delta = ctx->Input(2);
       xla::XlaOp limit = ctx->Input(1);
       xla::XlaOp start = ctx->Input(0);
-      if (type == DT_INT32 || type == DT_INT64) {
+      if (DataTypeIsInteger(type)) {
         auto dynamic_size = (xla::Abs(limit - start) + xla::Abs(delta) -
                              xla::One(ctx->builder(), ctx->input_xla_type(0))) /
                             xla::Abs(delta);

tensorflow/core/kernels/data/BUILD

@@ -95,6 +95,7 @@ tf_cc_test(
     deps = [
         ":cache_dataset_ops",
         ":iterator_ops",
+        ":range_dataset_op",
         ":tensor_slice_dataset_op",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",

tensorflow/core/kernels/data/cache_dataset_ops.cc

@@ -75,9 +75,15 @@ constexpr char kCacheDataset[] = "CacheDataset";
 constexpr char kIncompleteCacheErrorMessage[] =
     "The calling iterator did not fully read the dataset being cached. In "
     "order to avoid unexpected truncation of the dataset, the partially cached "
-    "contents of the dataset  will be discarded. This can happen if you have "
-    "an input pipeline similar to `dataset.cache().take(k).repeat()`. You "
-    "should use `dataset.take(k).cache().repeat()` instead.";
+    "contents of the dataset will be discarded. This can happen if you have "
+    "an input pipeline similar to `dataset.cache().take(k).repeat()`, or if "
+    "downstream operations drop elements (e.g. `batch(drop_remainder=True)`). "
+    "You should use `dataset.take(k).cache().repeat()` instead, or ensure the "
+    "dataset size is a multiple of the batch size before caching. Another "
+    "common workaround is to place the `.cache()` operation after the "
+    "operation that drops elements (like `.batch(...)`), if caching the "
+    "transformed data is acceptable.";
+constexpr size_t kMaxItems = 10000000;  // 10 million
 }  // namespace
 
 class DatasetRandomAccessCache {
@@ -89,15 +95,15 @@ class DatasetRandomAccessCache {
   // out_tensors with the element at that index.
   absl::Status Get(OpKernelContext* ctx, int64_t index,
                    std::vector<Tensor>* out_tensors) {
+    if (index < 0) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("Expected index >= 0; Received index: ", index));
+    }
     if (!iter_resource_) {
       TF_ASSIGN_OR_RETURN(iter_resource_,
                           GetIteratorResourceFromDataset(ctx, input_));
       TF_RETURN_IF_ERROR(iter_resource_->SetIteratorFromDataset(ctx, input_));
     }
-    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));
     }
@@ -159,12 +165,25 @@ class IteratorRandomAccessCache {
                        element_position));
     }
 
+    if (static_cast<size_t>(element_position) ==
+            std::numeric_limits<size_t>::max() ||
+        static_cast<size_t>(element_position) >= cache_.max_size()) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("Element position too large or invalid."));
+    }
+
     if (element_position < static_cast<int64_t>(cache_.size()) &&
         !cache_[element_position].empty()) {
       *out_tensors = cache_[element_position];
       return absl::OkStatus();
     }
 
+    if (element_position >= kMaxItems) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "Requested element_position ", element_position,
+          " exceeds the maximum allowed cache size of ", kMaxItems));
+    }
+
     TF_RETURN_IF_ERROR(input_->Get(ctx, element_position, out_tensors));
     if (element_position >= static_cast<int64_t>(cache_.size())) {
       cache_.resize(element_position + 1);
@@ -721,7 +740,6 @@ class CacheDatasetOp::FileDatasetBase : public DatasetBase {
   Env* const env_;
   const size_t num_tensors_;
   const size_t tensor_index_padding_size_;
-  static constexpr size_t kMaxItems = 10000000;  // 10 million
   const size_t item_index_padding_size_;
 };  // FileDatasetBase
 

tensorflow/core/kernels/data/cache_dataset_ops_test.cc

@@ -375,14 +375,51 @@ INSTANTIATE_TEST_CASE_P(CacheDatasetOpTest,
                         ParameterizedIteratorSaveAndRestoreTest,
                         ::testing::ValuesIn(IteratorSaveAndRestoreTestCases()));
 
-TEST_F(CacheDatasetOpTest, NegativeIndexTest) {
-  auto params = CacheDatasetParams3();
+TEST_F(CacheDatasetOpTest, NegativeIndexEarlyRejection) {
+  auto range_dataset_params = RangeDatasetParams(0, 20000000, 1);
+  auto params =
+      CacheDatasetParams(range_dataset_params,
+                         /*filename=*/"",
+                         /*output_dtypes=*/{DT_INT64},
+                         /*output_shapes=*/{PartialTensorShape({})}, kNodeName);
   TF_ASSERT_OK(Initialize(params));
   std::vector<Tensor> out_tensors;
   absl::Status status =
-      dataset_->Get(AnyContext(iterator_ctx_.get()), -1, &out_tensors);
-  EXPECT_TRUE(status.code() == absl::StatusCode::kOutOfRange);
-  EXPECT_EQ(status.message(), "Index out of range [0, 3):-1");
+      dataset_->Get(AnyContext(iterator_ctx_.get()), -1LL, &out_tensors);
+  EXPECT_TRUE(status.code() == absl::StatusCode::kInvalidArgument ||
+              status.code() == absl::StatusCode::kOutOfRange);
+}
+
+TEST_F(CacheDatasetOpTest, LargeIndexTest) {
+  auto range_dataset_params = RangeDatasetParams(0, 20000000, 1);
+  auto params =
+      CacheDatasetParams(range_dataset_params,
+                         /*filename=*/"",
+                         /*output_dtypes=*/{DT_INT64},
+                         /*output_shapes=*/{PartialTensorShape({})}, kNodeName);
+  TF_ASSERT_OK(Initialize(params));
+  std::vector<Tensor> out_tensors;
+  int64_t huge_index = std::numeric_limits<int64_t>::max();
+  absl::Status status =
+      dataset_->Get(AnyContext(iterator_ctx_.get()), huge_index, &out_tensors);
+  EXPECT_TRUE(status.code() == absl::StatusCode::kInvalidArgument ||
+              status.code() == absl::StatusCode::kOutOfRange);
+}
+
+TEST_F(CacheDatasetOpTest, BadAllocCrashTest) {
+  auto range_dataset_params = RangeDatasetParams(0, 20000000, 1);
+  auto params =
+      CacheDatasetParams(range_dataset_params,
+                         /*filename=*/"",
+                         /*output_dtypes=*/{DT_INT64},
+                         /*output_shapes=*/{PartialTensorShape({})}, kNodeName);
+  TF_ASSERT_OK(Initialize(params));
+  std::vector<Tensor> out_tensors;
+  absl::Status status =
+      dataset_->Get(AnyContext(iterator_ctx_.get()), 15000000, &out_tensors);
+  EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument);
+  EXPECT_TRUE(absl::StrContains(status.message(),
+                                "exceeds the maximum allowed cache size"));
 }
 
 }  // namespace

tensorflow/core/kernels/dilation_ops.cc

@@ -53,6 +53,9 @@ void ParseAttributes(OpKernelConstruction* context,
   OP_REQUIRES(context, (*strides)[0] == 1 && (*strides)[3] == 1,
               absl::UnimplementedError(
                   "Stride is only supported across spatial dimensions."));
+  OP_REQUIRES(context, (*strides)[1] >= 1 && (*strides)[2] >= 1,
+              absl::InvalidArgumentError(
+                  "Strides in the spatial dimensions must be >= 1."));
 
   OP_REQUIRES_OK(context, context->GetAttr("rates", rates));
   OP_REQUIRES(context, rates->size() == 4,
@@ -61,6 +64,9 @@ void ParseAttributes(OpKernelConstruction* context,
   OP_REQUIRES(context, (*rates)[0] == 1 && (*rates)[3] == 1,
               absl::UnimplementedError(
                   "Rate is only supported across spatial dimensions."));
+  OP_REQUIRES(context, (*rates)[1] >= 1 && (*rates)[2] >= 1,
+              absl::InvalidArgumentError(
+                  "Rates in the spatial dimensions must be >= 1."));
 
   OP_REQUIRES_OK(context, context->GetAttr("padding", padding));
 }
@@ -103,10 +109,12 @@ void ParseSizes(OpKernelContext* context, const std::vector<int32_t>& strides,
 
   // Effective filter size, after introducing rate - 1 zeros between each
   // non-zero filter element.
-  const int filter_rows_eff =
-      filter_rows + (filter_rows - 1) * (*rate_rows - 1);
-  const int filter_cols_eff =
-      filter_cols + (filter_cols - 1) * (*rate_cols - 1);
+  const int64_t filter_rows_eff =
+      static_cast<int64_t>(filter_rows) +
+      static_cast<int64_t>(filter_rows - 1) * (*rate_rows - 1);
+  const int64_t filter_cols_eff =
+      static_cast<int64_t>(filter_cols) +
+      static_cast<int64_t>(filter_cols - 1) * (*rate_cols - 1);
 
   OP_REQUIRES_OK(context, GetWindowedOutputSize(
                               input_rows, filter_rows_eff, /*dilation_rate=*/1,

tensorflow/core/kernels/map_stage_op.cc

@@ -448,11 +448,11 @@ class StagingMap : public ResourceBase {
 
     auto it = map_.begin();
 
+    *key = it->first;
+
     TF_RETURN_IF_ERROR(
         copy_or_move_tensors(&it->second, *key, *indices, tuple));
 
-    *key = it->first;
-
     // Remove entry if all the values have been consumed
     if (!std::any_of(
             it->second.begin(), it->second.end(),

tensorflow/core/profiler/utils/BUILD

@@ -261,7 +261,6 @@ cc_library(
     hdrs = ["hlo_module_utils.h"],
     visibility = internal_visibility([
         "//third_party/odml/model_explorer/backend/adapters/hlo:__pkg__",
-        "//tensorflow/compiler/mlir/lite/experimental/google/tooling/hlo_adapter:__pkg__",
     ]),
     deps = [
         "@org_xprof//xprof/utils:hlo_module_utils",

tensorflow/python/kernel_tests/data_structures/map_stage_op_test.py

@@ -655,6 +655,67 @@ def testNonScalarKeyMapUnStage(self):
       )
       self.evaluate(v)
 
+  def testMapUnstageNoKeyOutOfRangeIndex(self):
+    # MapUnstageNoKey with an out-of-range index after MapStage must surface
+    # a normal InvalidArgumentError from check_index(), not a fatal CHECK
+    # inside Tensor::CheckIsAlignedAndSingleElement. The CHECK can fire if
+    # popitem() reaches copy_or_move_tensors() with an empty local key
+    # tensor, because check_index() formats its error using
+    # key.scalar<int64_t>()() and Tensor::scalar() requires a 1-element
+    # tensor.
+    stage_op = data_flow_ops.gen_data_flow_ops.map_stage(
+        key=constant_op.constant([1], dtype=dtypes.int64),
+        indices=constant_op.constant([0], dtype=dtypes.int32),
+        values=[constant_op.constant([1.0], dtype=dtypes.float32)],
+        dtypes=[dtypes.float32],
+        capacity=10,
+        memory_limit=0,
+        container='',
+        shared_name='test_map_unstage_no_key_oob',
+        name=None,
+    )
+    self.evaluate(stage_op)
+    with self.assertRaisesRegex(errors.InvalidArgumentError, 'out of bounds'):
+      result = data_flow_ops.gen_data_flow_ops.map_unstage_no_key(
+          indices=[1],
+          dtypes=[dtypes.int64, dtypes.float32],
+          capacity=10,
+          memory_limit=0,
+          container='',
+          shared_name='test_map_unstage_no_key_oob',
+          name=None,
+      )
+      self.evaluate(result)
+
+  def testOrderedMapUnstageNoKeyOutOfRangeIndex(self):
+    # Parallel coverage for the ordered variant. OrderedMapUnstageNoKey
+    # shares StagingMap::popitem() with MapUnstageNoKey via the
+    # StagingMap<bool Ordered> template, so the same out-of-range-index
+    # path must surface InvalidArgumentError rather than abort.
+    stage_op = data_flow_ops.gen_data_flow_ops.ordered_map_stage(
+        key=constant_op.constant([1], dtype=dtypes.int64),
+        indices=constant_op.constant([0], dtype=dtypes.int32),
+        values=[constant_op.constant([1.0], dtype=dtypes.float32)],
+        dtypes=[dtypes.float32],
+        capacity=10,
+        memory_limit=0,
+        container='',
+        shared_name='test_ordered_map_unstage_no_key_oob',
+        name=None,
+    )
+    self.evaluate(stage_op)
+    with self.assertRaisesRegex(errors.InvalidArgumentError, 'out of bounds'):
+      result = data_flow_ops.gen_data_flow_ops.ordered_map_unstage_no_key(
+          indices=[1],
+          dtypes=[dtypes.int64, dtypes.float32],
+          capacity=10,
+          memory_limit=0,
+          container='',
+          shared_name='test_ordered_map_unstage_no_key_oob',
+          name=None,
+      )
+      self.evaluate(result)
+
 
 if __name__ == '__main__':
   test.main()