code: def test_gradient_norms_on_various_models(
      self,
      x_inputs,
      model_name,
      output_dim,
      layer_registry_name,
      per_example_loss_fn,
      num_microbatches,
      is_eager,
      partial,
  ):
      # Parse inputs to generate test data.
      is_ragged, input_data = x_inputs
      embed_indices = (
          tf.ragged.constant(input_data, dtype=tf.int32)
          if is_ragged
          else tf.convert_to_tensor(input_data, dtype_hint=tf.int32)
      )
  
      # The following are invalid test combinations and, hence, are skipped.
      batch_size = embed_indices.shape[0]
      if (
          (
              num_microbatches is not None
              and batch_size % num_microbatches != 0
          )
          or (model_name == "weighted_bow1" and is_ragged)
          or (
              # Current clipping ops do not have corresponding TPU kernels.
              self.using_tpu
              and is_ragged
          )
      ):
          return
  
      # Load shared assets to all devices.
      with self.strategy.scope():
  
          def embed_layer_generator(_, output_dims):
              return tf.keras.layers.Embedding(10, *output_dims)
  
          model = common_test_utils.get_model_from_generator(
              model_generator=get_embedding_model_generators()[
                  model_name
              ],
              layer_generator=embed_layer_generator,
              input_dims=embed_indices.shape[1:],
              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):
          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,
                  registry=get_embedding_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, autograph=False)
  
      # Set up the device ops and run the test.
      computed_norms, true_norms = self.strategy.run(
          test_op, args=(embed_indices,)
      )
      # 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=1e-3, atol=1e-2
      )
  
file_path: privacy/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/embedding_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
  from tensorflow_privacy.privacy.fast_gradient_clipping.registry_functions import embedding
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(x_inputs, model_name, output_dim, layer_registry_name, per_example_loss_fn, num_microbatches, is_eager, partial) and returns: common_test_utils, tf