Disable TF32 in a few tests.

tensorflow/c/eager/BUILD

@@ -427,6 +427,7 @@ tf_cuda_cc_test(
         "//tensorflow/c:tf_status_helper",
         "//tensorflow/c:tf_tensor_internal",
         "//tensorflow/c/experimental/ops",
+        "//tensorflow/core/platform:tensor_float_32_utils",
         "//tensorflow/core:test",
         "//tensorflow/core:test_main",
         "@com_google_absl//absl/types:span",

tensorflow/c/eager/gradient_checker_test.cc

@@ -20,6 +20,7 @@ limitations under the License.
 #include "tensorflow/c/experimental/ops/math_ops.h"
 #include "tensorflow/c/tf_status_helper.h"
 #include "tensorflow/c/tf_tensor.h"
+#include "tensorflow/core/platform/tensor_float_32_utils.h"
 #include "tensorflow/core/platform/test.h"
 
 namespace tensorflow {
@@ -96,6 +97,11 @@ class GradientCheckerTest
       ASSERT_EQ(errors::OK, s.code()) << s.error_message();
       ctx_.reset(ctx_raw);
     }
+
+    // Computing numerical gradients with TensorFloat-32 is numerically
+    // unstable. Some forward pass tests also fail with TensorFloat-32 due to
+    // low tolerances
+    enable_tensor_float_32_execution(false);
   }
 
   AbstractContextPtr ctx_;

tensorflow/compiler/xla/client/lib/BUILD

@@ -537,6 +537,7 @@ xla_test(
         "//tensorflow/compiler/xla/tests:literal_test_util",
         "//tensorflow/compiler/xla/tests:test_macros_header",
         "//tensorflow/compiler/xla/tests:xla_internal_test_main",
+        "//tensorflow/core/platform:tensor_float_32_utils",
         "//tensorflow/core:test",
     ],
 )

tensorflow/compiler/xla/client/lib/svd_test.cc

@@ -33,6 +33,7 @@ limitations under the License.
 #include "tensorflow/compiler/xla/tests/test_macros.h"
 #include "tensorflow/compiler/xla/xla_data.pb.h"
 #include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/tensor_float_32_utils.h"
 
 namespace xla {
 
@@ -54,6 +55,9 @@ class SVDTest : public ClientLibraryTestBase {
             {12, 48, 6, 62, 3},
         },
     };
+
+    // Test fails with TensorFloat-32 enabled
+    tensorflow::enable_tensor_float_32_execution(false);
   }
   void TearDown() override { ClientLibraryTestBase::TearDown(); }
 

tensorflow/python/keras/layers/gru_v2_test.py

@@ -59,6 +59,8 @@ _graph_options = config_pb2.GraphOptions(rewrite_options=_rewrites)
 _config = config_pb2.ConfigProto(graph_options=_graph_options)
 
 
+@tf_test_util.run_all_without_tensor_float_32(
+    "RNN GRU can use TF32 on GPU")
 @keras_parameterized.run_all_keras_modes(config=_config)
 class GRUV2Test(keras_parameterized.TestCase):
 
@@ -667,6 +669,8 @@ class GRUV2Test(keras_parameterized.TestCase):
     self.assertAllClose(self.evaluate(outputs), self.evaluate(copied_outputs))
 
 
+@tf_test_util.run_all_without_tensor_float_32(
+    "RNN GRU can use TF32 on GPU")
 class GRULayerGradientTapeTest(keras_parameterized.TestCase):
 
   @combinations.generate(combinations.combine(mode=['eager']))
@@ -694,6 +698,8 @@ class GRULayerGradientTapeTest(keras_parameterized.TestCase):
       tape.gradient(loss, gru.variables)
 
 
+@tf_test_util.run_all_without_tensor_float_32(
+    "RNN GRU can use TF32 on GPU")
 @keras_parameterized.run_all_keras_modes(config=_config)
 class GRUGraphRewriteTest(keras_parameterized.TestCase):