Simplify the structure of embedding, layer_normalization, and multi_head_attention tests.

PiperOrigin-RevId: 908775823

Author: ZacharyGarrett

tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/BUILD

@@ -85,6 +85,9 @@ py_library(
 py_test(
     name = "embedding_test",
     srcs = ["embedding_test.py"],
+    env = {
+        "TESTBRIDGE_TEST_ONLY": "GradNormTest.",
+    },
     shard_count = 12,
     deps = [
         ":dense",
@@ -107,6 +110,9 @@ py_library(
 py_test(
     name = "nlp_on_device_embedding_test",
     srcs = ["nlp_on_device_embedding_test.py"],
+    env = {
+        "TESTBRIDGE_TEST_ONLY": "GradNormTest.",
+    },
     shard_count = 6,
     deps = [
         ":dense",
@@ -129,6 +135,9 @@ py_library(
 py_test(
     name = "nlp_position_embedding_test",
     srcs = ["nlp_position_embedding_test.py"],
+    env = {
+        "TESTBRIDGE_TEST_ONLY": "GradNormTest.",
+    },
     shard_count = 6,
     deps = [
         ":dense",
@@ -151,6 +160,9 @@ py_library(
 py_test(
     name = "layer_normalization_test",
     srcs = ["layer_normalization_test.py"],
+    env = {
+        "TESTBRIDGE_TEST_ONLY": "GradNormTest.",
+    },
     shard_count = 8,
     deps = [
         ":dense",
@@ -174,6 +186,9 @@ py_library(
 py_test(
     name = "multi_head_attention_test",
     srcs = ["multi_head_attention_test.py"],
+    env = {
+        "TESTBRIDGE_TEST_ONLY": "GradNormTest.",
+    },
     shard_count = 8,
     deps = [
         ":dense",

tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/embedding_test.py

@@ -151,5 +151,15 @@ def test_op(x_batch):
     self.assertAllClose(computed_norms, true_norms, rtol=1e-3, atol=1e-2)
 
 
+class GradNormTpuTest(GradNormTest):
+
+  def setUp(self):
+    tf.config.experimental.disable_mlir_bridge()
+    super().setUp()
+    self.strategy = common_test_utils.create_tpu_strategy()
+    self.assertIn('TPU', self.strategy.extended.worker_devices[0])
+    self.using_tpu = True
+
+
 if __name__ == '__main__':
   tf.test.main()

tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/embedding_tpu_test.py

@@ -153,5 +153,14 @@ def test_op(x_batch):
     self.assertAllClose(computed_norms, true_norms, rtol=rtol, atol=atol)
 
 
+class GradNormTpuTest(GradNormTest):
+
+  def setUp(self):
+    super().setUp()
+    self.strategy = common_test_utils.create_tpu_strategy()
+    self.assertIn('TPU', self.strategy.extended.worker_devices[0])
+    self.using_tpu = True
+
+
 if __name__ == '__main__':
   tf.test.main()

tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/layer_normalization_test.py

@@ -371,5 +371,14 @@ def test_op(query_batch, value_batch, key_batch, mask_batch):
     self.assertAllClose(computed_norms, true_norms, rtol=rtol, atol=atol)
 
 
+class GradNormTpuTest(GradNormTest):
+
+  def setUp(self):
+    super().setUp()
+    self.strategy = common_test_utils.create_tpu_strategy()
+    self.assertIn('TPU', self.strategy.extended.worker_devices[0])
+    self.using_tpu = True
+
+
 if __name__ == '__main__':
   tf.test.main()

tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/layer_normalization_tpu_test.py

@@ -145,5 +145,15 @@ def test_op(x_batch):
     self.assertAllClose(computed_norms, true_norms, rtol=1e-3, atol=1e-2)
 
 
+class GradNormTpuTest(GradNormTest):
+
+  def setUp(self):
+    tf.config.experimental.disable_mlir_bridge()
+    super().setUp()
+    self.strategy = common_test_utils.create_tpu_strategy()
+    self.assertIn('TPU', self.strategy.extended.worker_devices[0])
+    self.using_tpu = True
+
+
 if __name__ == '__main__':
   tf.test.main()

tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/multi_head_attention_test.py

@@ -136,5 +136,14 @@ def test_op():
     self.assertAllClose(computed_norms, true_norms, rtol=1e-3, atol=1e-2)
 
 
+class GradNormTpuTest(GradNormTest):
+
+  def setUp(self):
+    super().setUp()
+    self.strategy = common_test_utils.create_tpu_strategy()
+    self.assertIn('TPU', self.strategy.extended.worker_devices[0])
+    self.using_tpu = True
+
+
 if __name__ == '__main__':
   tf.test.main()