Disable TF32 for tests that fail with it

sonnet/src/conformance/xla_test.py

@@ -27,6 +27,10 @@ class XLATest(test_utils.TestCase, parameterized.TestCase):
 
   @goldens.all_goldens
   def test_compile(self, golden):
+    # Fails on GPUs where TF32 is enabled by default.
+    tf.config.experimental.enable_tensor_float_32_execution(False)
+    self.addCleanup(tf.config.experimental.enable_tensor_float_32_execution,
+                    True)
     mod = golden.create_module()
     golden.create_all_variables(mod)
 
@@ -55,6 +59,10 @@ def forward():
 
   @goldens.all_goldens
   def test_jit_scope(self, golden):
+    # Fails on GPUs where TF32 is enabled by default.
+    tf.config.experimental.enable_tensor_float_32_execution(False)
+    self.addCleanup(tf.config.experimental.enable_tensor_float_32_execution,
+                    True)
     mod = golden.create_module()
     golden.create_all_variables(mod)
 

sonnet/src/depthwise_conv_test.py

@@ -67,6 +67,11 @@ def testDefaultInitializers(self, dtype):
                                   ("ValidPaddingNoBias", False, "VALID"),
                                   ("ValidPaddingUseBias", True, "VALID"))
   def testFunction(self, with_bias, padding):
+    if self.primary_device == "GPU":
+      # Fails on GPUs where TF32 is enabled by default.
+      tf.config.experimental.enable_tensor_float_32_execution(False)
+      self.addCleanup(tf.config.experimental.enable_tensor_float_32_execution,
+                      True)
     conv1 = depthwise_conv.DepthwiseConv2D(
         channel_multiplier=1,
         kernel_shape=3,

sonnet/src/nets/dnc/util_test.py

@@ -114,6 +114,10 @@ def testCorrectOutput(self):
 class LinearTest(test_utils.TestCase, parameterized.TestCase):
 
   def testLinearOutputOneModule(self):
+    # Fails on GPUs where TF32 is enabled by default.
+    tf.config.experimental.enable_tensor_float_32_execution(False)
+    self.addCleanup(tf.config.experimental.enable_tensor_float_32_execution,
+                    True)
     batch_size = 4
     input_size = 5
     output_size = 3
@@ -126,6 +130,10 @@ def testLinearOutputOneModule(self):
     self.assertAllClose(expected_output, output.numpy(), atol=self.get_atol())
 
   def testLinearOutputTwoModules(self):
+    # Fails on GPUs where TF32 is enabled by default.
+    tf.config.experimental.enable_tensor_float_32_execution(False)
+    self.addCleanup(tf.config.experimental.enable_tensor_float_32_execution,
+                    True)
     batch_size = 4
     input_size_a = 5
     input_size_b = 6

sonnet/src/pad_test.py

@@ -65,6 +65,10 @@ def test_padding_reverse_causal(self):
   @parameterized.parameters((1, 1, 1), (3, 1, 1), (1, 3, 1), (1, 1, 3),
                             (3, 3, 1), (3, 1, 3), (1, 3, 3), (3, 3, 3))
   def test_same_padding(self, kernel_size, stride, rate):
+    # Fails on GPUs where TF32 is enabled by default.
+    tf.config.experimental.enable_tensor_float_32_execution(False)
+    self.addCleanup(tf.config.experimental.enable_tensor_float_32_execution,
+                    True)
     a = tf.random.normal([2, 4, 3])
     k = tf.random.normal([kernel_size, 3, 4])
     padding = pad.create(pad.same, kernel_size, rate, 1, -1)

sonnet/src/recurrent_test.py

@@ -36,6 +36,10 @@ def setUp(self):
 
   @parameterized.parameters([False, True])
   def testComputationAgainstNumPy(self, use_tf_function):
+    # Fails on GPUs where TF32 is enabled by default.
+    tf.config.experimental.enable_tensor_float_32_execution(False)
+    self.addCleanup(tf.config.experimental.enable_tensor_float_32_execution,
+                    True)
     inputs = self.evaluate(
         tf.random.uniform([self.batch_size, self.input_size]))
     core = recurrent.VanillaRNN(
@@ -164,6 +168,10 @@ def setUp(self):
       itertools.product([False, True], [None, 4], [0.0, 1.0]))
   def testComputationAgainstNumPy(self, use_tf_function, projection_size,
                                   forget_bias):
+    # Fails on GPUs where TF32 is enabled by default.
+    tf.config.experimental.enable_tensor_float_32_execution(False)
+    self.addCleanup(tf.config.experimental.enable_tensor_float_32_execution,
+                    True)
     inputs = self.evaluate(
         tf.random.uniform([self.batch_size, self.input_size]))
     core = recurrent.LSTM(
@@ -262,6 +270,10 @@ def setUp(self):
 
   @parameterized.parameters(itertools.product([1, 4], [True, False]))
   def testComputationAgainstLSTM(self, num_steps, use_tf_function):
+    # Fails on GPUs where TF32 is enabled by default.
+    tf.config.experimental.enable_tensor_float_32_execution(False)
+    self.addCleanup(tf.config.experimental.enable_tensor_float_32_execution,
+                    True)
     unrolled_lstm = recurrent.UnrolledLSTM(self.hidden_size)
     initial_state = unrolled_lstm.initial_state(self.batch_size)
 
@@ -449,6 +461,10 @@ def setUp(self):
 
   @parameterized.parameters([False, True])
   def testComputationAgainstNumPy(self, use_tf_function):
+    # Fails on GPUs where TF32 is enabled by default.
+    tf.config.experimental.enable_tensor_float_32_execution(False)
+    self.addCleanup(tf.config.experimental.enable_tensor_float_32_execution,
+                    True)
     inputs = self.evaluate(
         tf.random.uniform([self.batch_size, self.input_size]))
     core = recurrent.GRU(self.hidden_size)