Simplify the structure of einsum_dense_test.

PiperOrigin-RevId: 908333752

Author: ZacharyGarrett

tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/BUILD

@@ -35,6 +35,9 @@ py_test(
     name = "einsum_dense_test",
     size = "large",
     srcs = ["einsum_dense_test.py"],
+    env = {
+        "TESTBRIDGE_TEST_ONLY": "GradNormTest.",
+    },
     shard_count = 12,
     deps = [
         ":dense",

tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/einsum_dense_test.py

@@ -167,5 +167,14 @@ def test_op(x):
     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()