code: def test_gradient_norms_on_various_models(
      self,
      model_name,
      layer_name,
      layer_registry_name,
      param_tuple,
      num_microbatches,
      is_eager,
  ):
      # Parse inputs to generate test data.
      (
          query_input_dims,
          value_input_dims,
          use_key,
          use_attention_mask,
          num_heads,
          key_dim,
          value_dim,
          dropout,
          use_bias,
          output_shape,
      ) = param_tuple
      attention_generator_inputs = (
          num_heads,
          key_dim,
          value_dim,
          dropout,
          use_bias,
          output_shape,
      )
      (
          query_batch,
          value_batch,
          key_batch,
          mask_batch,
      ) = get_multi_head_attention_example_inputs(
          query_input_dims, value_input_dims
      )
      mask_input_dims = query_input_dims[:-1] + value_input_dims[:-1]
  
      # Make the layer generator via currying.
      attention_generator = get_attention_layer_generators()[layer_name]
  
      def curried_generator(a, b):
          del a, b
          return attention_generator(*attention_generator_inputs)
  
      # Load shared assets to all devices.
      with self.strategy.scope():
          model = common_test_utils.get_model_from_generator(
              model_generator=get_attention_model_generators()[
                  model_name
              ],
              layer_generator=curried_generator,
              input_dims=(
                  query_input_dims,
                  value_input_dims,
                  value_input_dims,
                  mask_input_dims,
                  use_key,
                  use_attention_mask,
              ),
              output_dims=None,
              is_eager=is_eager,
          )
  
      # Define the main testing ops. These may be later compiled to a Graph op.
      def test_op(query_batch, value_batch, key_batch, mask_batch):
          return (
              common_test_utils.get_computed_and_true_norms_from_model(
                  model=model,
                  per_example_loss_fn=None,
                  num_microbatches=num_microbatches,
                  x_batch=(
                      query_batch,
                      value_batch,
                      key_batch,
                      mask_batch,
                  ),
                  weight_batch=None,
                  registry=get_attention_layer_registries()[
                      layer_registry_name
                  ],
                  partial=False,
              )
          )
  
      # 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.756414
      #   true_norm     = 23.338600
      #   abs_diff      = 0.58218575
      #   rel_diff      = 0.02494519
      #
      # 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-1 if self.using_tpu else 1e-2
      atol = 1e-0 if self.using_tpu else 1e-1
  
      # Set up the device ops and run the test.
      computed_norms, true_norms = self.strategy.run(
          test_op,
          args=(query_batch, value_batch, key_batch, mask_batch),
      )
      # TPUs return replica contexts, which must be unwrapped.
      batch_size = tf.shape(query_batch)[0]
      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
      )
  
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
  from tensorflow_privacy.privacy.fast_gradient_clipping.registry_functions import multi_head_attention
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, layer_registry_name, param_tuple, num_microbatches, is_eager) and returns: attention_generator, common_test_utils