Merge pull request #52765 from pranve/cherrypick-9294094df6fea79271778eb7e7ae1bad8b5ef98f-on-r2.6

mihaimaruseac · web-flow · commit d27fbb48f2ae · 2021-10-28T15:10:47.000-07:00
Merge pull request #51138 from yongtang:46911-tile
diff --git a/tensorflow/core/kernels/tile_ops.cc b/tensorflow/core/kernels/tile_ops.cc
@@ -188,7 +188,8 @@ class TileOp : public OpKernel {
           context, multiples_array[i] >= 0,
           errors::InvalidArgument("Expected multiples[", i, "] >= 0, but got ",
                                   multiples_array[i]));
-      output_shape.AddDim(input.dim_size(i) * multiples_array[i]);
+      OP_REQUIRES_OK(context, output_shape.AddDimWithStatus(
+                                  input.dim_size(i) * multiples_array[i]));
     }
     if (output_shape == input.shape()) {
       context->set_output(0, input);
diff --git a/tensorflow/python/kernel_tests/shape_ops_test.py b/tensorflow/python/kernel_tests/shape_ops_test.py
@@ -723,6 +723,17 @@ def testShapeFunctionEdgeCases(self):
           inp, array_ops.placeholder(
               dtypes.int32, shape=[3]))
 
+  def testLargeTensor(self):
+    # Test case for GItHub issue 46911.
+    if test_util.is_xla_enabled():
+      # The following test fails with XLA enabled.
+      return
+    with self.assertRaises(errors_impl.InternalError):
+      with self.cached_session():
+        tiled = array_ops.tile(
+            np.ones((1, 1, 1)), [100000000, 100000000, 100000000])
+        self.evaluate(tiled)
+
 
 if __name__ == "__main__":
   test.main()
