diff --git a/sonnet/src/conformance/xla_test.py b/sonnet/src/conformance/xla_test.py index bf5e18f..339bdbf 100644 --- a/sonnet/src/conformance/xla_test.py +++ b/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) diff --git a/sonnet/src/depthwise_conv_test.py b/sonnet/src/depthwise_conv_test.py index 1c4215e..15df5ef 100644 --- a/sonnet/src/depthwise_conv_test.py +++ b/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, diff --git a/sonnet/src/nets/dnc/util_test.py b/sonnet/src/nets/dnc/util_test.py index a7ed9f4..5ffef45 100644 --- a/sonnet/src/nets/dnc/util_test.py +++ b/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 diff --git a/sonnet/src/pad_test.py b/sonnet/src/pad_test.py index ba42104..29e3e2e 100644 --- a/sonnet/src/pad_test.py +++ b/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) diff --git a/sonnet/src/recurrent_test.py b/sonnet/src/recurrent_test.py index 415346d..8a9c3ae 100644 --- a/sonnet/src/recurrent_test.py +++ b/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)