code: def test_gradient_norms_on_various_models(
      self,
      model_name,
      layer_name,
      input_dims,
      output_dim,
      layer_registry_name,
      per_example_loss_fn,
      num_microbatches,
      is_eager,
      partial,
      weighted,
  ):
      # Parse inputs to generate test data.
      x_batches, weight_batches = common_test_utils.get_nd_test_batches(
          input_dims[0]
      )
  
      # Load shared assets to all devices.
      with self.strategy.scope():
          model = common_test_utils.get_model_from_generator(
              model_generator=get_dense_model_generators()[model_name],
              layer_generator=get_dense_layer_generators()[layer_name],
              input_dims=input_dims,
              output_dims=[output_dim],
              is_eager=is_eager,
          )
  
      # Define the main testing ops. These may be later compiled to a Graph op.
      def test_op(x_batch, weight_batch):
          return (
              common_test_utils.get_computed_and_true_norms_from_model(
                  model=model,
                  per_example_loss_fn=per_example_loss_fn,
                  num_microbatches=num_microbatches,
                  x_batch=[x_batch, x_batch]
                  if model_name == "tower2"
                  else x_batch,
                  weight_batch=weight_batch if weighted else None,
                  registry=get_dense_layer_registries()[
                      layer_registry_name
                  ],
                  partial=partial,
              )
          )
  
      # TPUs can only run `tf.function`-decorated functions.
      if self.using_tpu:
          test_op = tf.function(
              test_op, jit_compile=True, autograph=False
          )
  
      # TPUs use lower precision than CPUs, so we relax our criterion.
      # E.g., one of the TPU runs generated the following results:
      #
      #   computed_norm = 22.530651
      #   true_norm     = 22.570976
      #   abs_diff      = 0.04032516
      #   rel_diff      = 0.00178659
      #
      # which is a reasonable level of error for computing gradient norms.
      # Other trials also give an absolute (resp. relative) error of around
      # 0.05 (resp. 0.0015).
      rtol = 1e-2 if self.using_tpu else 1e-3
      atol = 1e-1 if self.using_tpu else 1e-2
  
      for x_batch, weight_batch in zip(x_batches, weight_batches):
          batch_size = x_batch.shape[0]
          if (
              num_microbatches is not None
              and batch_size % num_microbatches != 0
          ):
              continue
          # Set up the device ops and run the test.
          computed_norms, true_norms = self.strategy.run(
              test_op, args=(x_batch, weight_batch)
          )
          # TPUs return replica contexts, which must be unwrapped.
          if self.using_tpu:
              common_test_utils.assert_replica_values_are_close(
                  self, computed_norms
              )
              common_test_utils.assert_replica_values_are_close(
                  self, true_norms
              )
              computed_norms = computed_norms.values[0]
              true_norms = true_norms.values[0]
          expected_size = num_microbatches or batch_size
          self.assertEqual(tf.shape(computed_norms)[0], expected_size)
          self.assertAllClose(
              computed_norms, true_norms, rtol=rtol, atol=atol
          )
  
file_path: privacy/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/dense_test.py
imports: from absl.testing import parameterized
  import tensorflow
  from tensorflow_privacy.privacy.fast_gradient_clipping import common_test_utils
  from tensorflow_privacy.privacy.fast_gradient_clipping import layer_registry
  from tensorflow_privacy.privacy.fast_gradient_clipping.registry_functions import dense
label_desc: Use image node_id 2 for calling the GradNormTest obj's underlying member method code with example usage: obj.test_gradient_norms_on_various_models(model_name, layer_name, input_dims, output_dim, layer_registry_name, per_example_loss_fn, num_microbatches, is_eager, partial, weighted) and returns: common_test_utils