Merge pull request #52765 from pranve/cherrypick-9294094df6fea7927177…

…8eb7e7ae1bad8b5ef98f-on-r2.6 Merge pull request #51138 from yongtang:46911-tile
tensorflow · Oct 28, 2021 · d27fbb4 · d27fbb4
2 parents dd68305 + cbe7645
commit d27fbb4
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
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()