Simplify the structure of the dense_test.

PiperOrigin-RevId: 895460903

Author: ZacharyGarrett

tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/BUILD

@@ -59,10 +59,12 @@ py_test(
     name = "dense_test",
     size = "large",
     srcs = ["dense_test.py"],
+    env = {
+        "TESTBRIDGE_TEST_ONLY": "GradNormTest.",
+    },
     shard_count = 12,
     deps = [
         ":dense",
-        "//tensorflow_privacy/privacy/fast_gradient_clipping:clip_grads",
         "//tensorflow_privacy/privacy/fast_gradient_clipping:common_test_utils",
         "//tensorflow_privacy/privacy/fast_gradient_clipping:layer_registry",
     ],

tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/dense_test.py

@@ -142,5 +142,14 @@ def test_op(x_batch, weight_batch):
       self.assertAllClose(computed_norms, true_norms, rtol=rtol, atol=atol)
 
 
+class GradNormTpuTest(GradNormTest):
+
+  def setUp(self):
+    super(GradNormTest, self).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()