Fixed mixed precision setting in test.

PiperOrigin-RevId: 592423175
Change-Id: I83bee20c5cda13beab8d1a9191d0f1bf7d178bf3

Author: ssjhv

tensorflow_compression/python/entropy_models/continuous_batched_test.py

@@ -215,7 +215,7 @@ def test_dtypes_are_correct_with_mixed_precision(self):
       self.assertEqual(bits.shape, (2,))
       self.assertAllGreaterEqual(bits, 0.)
     finally:
-      tf.keras.mixed_precision.set_global_policy(None)
+      tf.keras.mixed_precision.set_global_policy(tf.keras.backend.floatx())
 
   def test_small_cdfs_for_dirac_prior_without_quantization_offset(self):
     prior = uniform_noise.NoisyNormal(loc=100. * tf.range(16.), scale=1e-10)

tensorflow_compression/python/entropy_models/continuous_indexed_test.py

@@ -220,7 +220,7 @@ def test_dtypes_are_correct_with_mixed_precision(self):
       self.assertEqual(bits.shape, (2,))
       self.assertAllGreaterEqual(bits, 0.)
     finally:
-      tf.keras.mixed_precision.set_global_policy(None)
+      tf.keras.mixed_precision.set_global_policy(tf.keras.backend.floatx())
 
 
 class LocationScaleIndexedEntropyModelTest(tf.test.TestCase):

tensorflow_compression/python/entropy_models/power_law_test.py

@@ -115,7 +115,7 @@ def test_dtypes_are_correct_with_mixed_precision(self):
       self.assertEqual(penalty.dtype, tf.float16)
       self.assertEqual(penalty.shape, (2,))
     finally:
-      tf.keras.mixed_precision.set_global_policy(None)
+      tf.keras.mixed_precision.set_global_policy(tf.keras.backend.floatx())
 
 
 if __name__ == "__main__":

tensorflow_compression/python/layers/gdn_test.py

@@ -207,7 +207,7 @@ def test_dtypes_are_correct_with_mixed_precision(self):
         self.assertEqual(variable.dtype, tf.float32)
       self.assertEqual(y.dtype, tf.float16)
     finally:
-      tf.keras.mixed_precision.set_global_policy(None)
+      tf.keras.mixed_precision.set_global_policy(tf.keras.backend.floatx())
 
 
 if __name__ == "__main__":

tensorflow_compression/python/layers/signal_conv_test.py

@@ -162,7 +162,7 @@ def test_dtypes_are_correct_with_mixed_precision(self):
         self.assertEqual(variable.dtype, tf.float32)
       self.assertEqual(y.dtype, tf.float16)
     finally:
-      tf.keras.mixed_precision.set_global_policy(None)
+      tf.keras.mixed_precision.set_global_policy(tf.keras.backend.floatx())
 
 
 class ConvolutionsTest(tf.test.TestCase):