Disable MLIR bridge for the test points that MLIR bridge silently fails

PiperOrigin-RevId: 676573736
tensorflow · Sep 19, 2024 · 301c5c4 · 301c5c4
1 parent e885683
commit 301c5c4
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/embedding_tpu_test.py b/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/embedding_tpu_test.py
@@ -20,6 +20,7 @@
 class GradNormTpuTest(embedding_test.GradNormTest):
 
   def setUp(self):
+    tf.config.experimental.disable_mlir_bridge()
     super(embedding_test.GradNormTest, self).setUp()
     self.strategy = common_test_utils.create_tpu_strategy()
     self.assertIn('TPU', self.strategy.extended.worker_devices[0])

diff --git a/...acy/privacy/fast_gradient_clipping/registry_functions/nlp_on_device_embedding_tpu_test.py b/...acy/privacy/fast_gradient_clipping/registry_functions/nlp_on_device_embedding_tpu_test.py
@@ -20,6 +20,7 @@
 class GradNormTpuTest(nlp_on_device_embedding_test.GradNormTest):
 
   def setUp(self):
+    tf.config.experimental.disable_mlir_bridge()
     super(nlp_on_device_embedding_test.GradNormTest, self).setUp()
     self.strategy = common_test_utils.create_tpu_strategy()
     self.assertIn('TPU', self.strategy.extended.worker_devices[0])