Addresses precision issues in unit tests on A100 platform.

This should fix
#155

PiperOrigin-RevId: 487422439
Change-Id: I612ad5e2fdaed1e50bc36fcb17723bfa6ea4d0a1

Author: Johannes Ballé

tensorflow_compression/python/distributions/deep_factorized_test.py

@@ -23,6 +23,12 @@
 
 class DeepFactorizedTest(tf.test.TestCase, parameterized.TestCase):
 
+  def setUp(self):
+    super().setUp()
+    # Disable TensorFloat-32 format on A100 platform, as precision is too low
+    # for current test assertions.
+    tf.config.experimental.enable_tensor_float_32_execution(False)
+
   def test_can_instantiate_scalar(self):
     df = deep_factorized.DeepFactorized()
     self.assertEqual(df.batch_shape, ())
@@ -66,6 +72,12 @@ def test_broadcasts_correctly(self, method):
 
 class NoisyDeepFactorizedTest(tf.test.TestCase):
 
+  def setUp(self):
+    super().setUp()
+    # Disable TensorFloat-32 format on A100 platform, as precision is too low
+    # for current test assertions.
+    tf.config.experimental.enable_tensor_float_32_execution(False)
+
   def test_can_instantiate_and_run_scalar(self):
     df = deep_factorized.NoisyDeepFactorized(num_filters=(2, 3, 4))
     self.assertEqual(df.batch_shape, ())

tensorflow_compression/python/layers/gdn_test.py

@@ -23,6 +23,12 @@
 
 class GDNTest(tf.test.TestCase, parameterized.TestCase):
 
+  def setUp(self):
+    super().setUp()
+    # Disable TensorFloat-32 format on A100 platform, as precision is too low
+    # for current test assertions.
+    tf.config.experimental.enable_tensor_float_32_execution(False)
+
   def test_invalid_data_format_raises_error(self):
     with self.assertRaises(ValueError):
       gdn.GDN(data_format="NHWC")