code: def compute_true_gradient_norms(
      input_model: tf.keras.Model,
      x_batch: tf.Tensor,
      y_batch: tf.Tensor,
      weight_batch: Optional[tf.Tensor],
      per_example_loss_fn: Optional[type_aliases.LossFn],
      num_microbatches: Optional[int],
      trainable_vars: Optional[tf.Variable] = None,
  ) -> type_aliases.OutputTensors:
      """Computes the real gradient norms for an input `(model, x, y)`."""
      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
          )
      with tf.GradientTape(persistent=True) as tape:
          y_pred = input_model(x_batch)
          loss = per_example_loss_fn(y_batch, y_pred, weight_batch)
          if num_microbatches is not None:
              loss = tf.reduce_mean(
                  tf.reshape(
                      loss,
                      tf.concat(
                          [[num_microbatches, -1], tf.shape(loss)[1:]],
                          axis=0,
                      ),
                  ),
                  axis=1,
              )
          if isinstance(loss, tf.RaggedTensor):
              loss = loss.to_tensor()
      sqr_norms = []
      trainable_vars = trainable_vars or input_model.trainable_variables
      for var in trainable_vars:
          jacobian = tape.jacobian(
              loss, var, experimental_use_pfor=False
          )
          reduction_axes = tf.range(1, tf.rank(jacobian))
          sqr_norm = tf.reduce_sum(
              tf.square(jacobian), axis=reduction_axes
          )
          sqr_norms.append(sqr_norm)
      sqr_norm_tsr = tf.stack(sqr_norms, axis=1)
      return tf.sqrt(tf.reduce_sum(sqr_norm_tsr, axis=1))
  
imports: from collections.abc import Callable, MutableSequence, Sequence
  from typing import Optional
  import numpy
  import tensorflow
  import tensorflow.compat.v2
  from tensorflow_privacy.privacy.fast_gradient_clipping import clip_grads
  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 5 for calling a global function with example usage: compute_true_gradient_norms(input_model, x_batch, y_batch, weight_batch, per_example_loss_fn, num_microbatches, trainable_vars) and returns: tf