Disabling unit tests that fail with ROCm

tensorflow/python/keras/distribute/BUILD

@@ -425,6 +425,7 @@ distribute_py_test(
     shard_count = 8,
     tags = [
         "multi_and_single_gpu",
+        "no_rocm",
         "no_windows_gpu",
         "notsan",
     ],

tensorflow/python/keras/integration_test/gradient_checkpoint_test.py

@@ -16,6 +16,7 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+from tensorflow.python.platform import test as test_lib
 import gc
 
 import tensorflow as tf
@@ -150,6 +151,10 @@ class GradientCheckpointTest(tf.test.TestCase):
   def test_does_not_raise_oom_exception(self):
     if not _limit_gpu_memory():
       self.skipTest('No virtual GPUs found')
+    if test_lib.is_built_with_rocm():
+      self.skipTest('ROCm MIOpen does not support searching for memory-limited'
+                    'solvers yet so skip the subtest which would result in OOM.'
+                   )
     n_step = 2
     losses = _train_with_recompute(n_step)
     self.assertLen(losses, n_step)

tensorflow/python/keras/mixed_precision/experimental/BUILD

@@ -69,6 +69,7 @@ py_test(
     ],
     python_version = "PY3",
     srcs_version = "PY2AND3",
+    tags = ["no_rocm"],
     deps = [
         ":policy",
         "//tensorflow/python:client_testlib",

tensorflow/python/keras/mixed_precision/experimental/layer_correctness_test.py

@@ -159,6 +159,10 @@ class LayerCorrectnessTest(keras_parameterized.TestCase):
       input_data: A Numpy array with the data of the input. If None, input data
         will be randomly generated
     """
+
+    if f32_layer_fn == convolutional.ZeroPadding2D and \
+       test.is_built_with_rocm():
+      return
     if isinstance(input_shape[0], int):
       input_shapes = [input_shape]
     else: