code: def compute_gradient_norms(
      input_model: tf.keras.Model,
      layer_registry: lr.LayerRegistry,
      x_batch: type_aliases.InputTensors,
      y_batch: tf.Tensor,
      weight_batch: Optional[tf.Tensor] = None,
      per_example_loss_fn: Optional[type_aliases.LossFn] = None,
      num_microbatches: Optional[type_aliases.BatchSize] = None,
      trainable_vars: Optional[Sequence[tf.Variable]] = None,
  ):
      """Computes the per-example loss gradient norms for given data.
  
      Applies a variant of the approach given in
        https://arxiv.org/pdf/2009.03106.pdf
  
      Args:
        input_model: The `tf.keras.Model` from which to obtain the layers from. The
          loss of the model *must* be a scalar loss. When using microbatching, the
          loss reduction must be mean.
        layer_registry: A `LayerRegistry` instance containing functions that help
          compute gradient norms quickly. See
          `tensorflow_privacy.privacy.fast_gradient_clipping.layer_registry` for
          more details.
        x_batch: An `InputTensor` representing a batch of inputs to the model. The
          first axis must be the batch dimension.
        y_batch: A `tf.Tensor` representing a batch of output labels. The first axis
          must be the batch dimension. The number of examples should match the
          number of examples in `x_batch`.
        weight_batch: Optional batch of weights, passed to the loss function.
          Weights apply to the loss prior to clipping.
        per_example_loss_fn: takes as input predictions, labels and weights, and
          outputs a vector of per-example losses. If None, derived from
          `input_model.loss` by disabling its reduction.
        num_microbatches: An optional number or scalar `tf.Tensor` for the number of
          microbatches. If not None, indicates that the loss is grouped into
          num_microbatches (in this case, the batch dimension needs to be a multiple
          of num_microbatches). When there is microbatches, we always assume the
          loss is the mean over a microbatch. And the gradient norm is computed for
          each microbatch.
        trainable_vars: The list of variables included in computing the gradient
          norm. When a layer has multiple variables, we include all the variables if
          any of the variables is in the list. If `trainable_vars` is None, all the
          variables are included.
  
      Returns:
        A scalar vector, whose i-th entry is the norm of the gradient of the i-th
        weighted example loss (when num_microbatches is None) or the norm of the
        gradient of the i-th microbatch loss (define as a mean over the microbatch).
        Note that when the loss is weighted (`weight_batch` is not None), weights
        are applied prior to clipping.
      """
      tape = tf.GradientTape(
          persistent=True, watch_accessed_variables=False
      )
      registry_generator_fn = get_registry_generator_fn(
          tape, layer_registry, num_microbatches
      )
      # First loop computes the model outputs, summed loss, and generator outputs.
      with tape:
          (
              model_outputs,
              generator_outputs_list,
          ) = gradient_clipping_utils.model_forward_pass(
              input_model, x_batch, generator_fn=registry_generator_fn
          )
          # Ignore the original loss function's reduction to get per-example loss.
          if per_example_loss_fn is None:
              loss_config = input_model.loss.get_config()
              loss_config["reduction"] = tf.keras.losses.Reduction.NONE
              per_example_loss_fn = input_model.loss.from_config(
                  loss_config
              )
          losses = per_example_loss_fn(
              y_batch, model_outputs, weight_batch
          )
          if losses.shape is None:
              raise NotImplementedError(
                  "The unreduced (or per-example) loss's shape cannot be `None`"
              )
          if len(losses.shape) != 1:
              raise NotImplementedError(
                  "The unreduced (or per-example) loss needs to have a shape of length "
                  "one, but received an unreduced loss of shape length %s"
                  % len(losses.shape)
              )
          if num_microbatches is not None:
              losses = tf.reduce_mean(
                  common_manip_utils.maybe_add_microbatch_axis(
                      losses, num_microbatches
                  ),
                  axis=1,
              )
          summed_loss = tf.reduce_sum(losses)
      # Unwrap the generator outputs so that the next loop avoids duplicating
      # backprop ops.
      filtered_outputs = [
          t for t in generator_outputs_list if t is not None
      ]
      vars_list = []
      sqr_norm_fns_list = []
      if trainable_vars is not None:
          # Create a set using `ref()` for fast set membership check. tf.Variable
          # itself is not hashable.
          trainable_vars = set([v.ref() for v in trainable_vars])
      for v, f, weights_list in filtered_outputs:
          if trainable_vars is None or any(
              w.ref() in trainable_vars for w in weights_list
          ):
              # Include only those variables in trainable_vars.
              vars_list.append(v)
              sqr_norm_fns_list.append(f)
      # Second loop evaluates the squared L2 norm functions and appends the results.
      grads_list = tape.gradient(
          summed_loss,
          vars_list,
          unconnected_gradients=tf.UnconnectedGradients.ZERO,
      )
      if not grads_list:
          raise ValueError("The gradient list cannot be empty.")
      if len(grads_list) != len(sqr_norm_fns_list):
          raise ValueError("There must be as many norms as gradients.")
      sqr_norm_list = []
      for grads, f in zip(grads_list, sqr_norm_fns_list):
          sqr_norm_list.append(f(grads))
      sqr_norm_tsr = tf.stack(sqr_norm_list, axis=1)
      return tf.sqrt(tf.reduce_sum(sqr_norm_tsr, axis=1))
  
file_path: privacy/tensorflow_privacy/privacy/fast_gradient_clipping/clip_grads.py
imports: from collections.abc import Sequence
  from typing import Optional
  import tensorflow
  from tensorflow_privacy.privacy.fast_gradient_clipping import common_manip_utils
  from tensorflow_privacy.privacy.fast_gradient_clipping import gradient_clipping_utils
  from tensorflow_privacy.privacy.fast_gradient_clipping import layer_registry
  from tensorflow_privacy.privacy.fast_gradient_clipping import type_aliases
label_desc: Use image node_id 2 for calling a global function with example usage: compute_gradient_norms(input_model, layer_registry, x_batch, y_batch, weight_batch, per_example_loss_fn, num_microbatches, trainable_vars) and returns: tf