code: def gradient_penalty_loss(
      real_data: Union[tf.Tensor, Sequence[tf.Tensor]],
      generated_data: Union[tf.Tensor, Sequence[tf.Tensor]],
      discriminator: tf.keras.Model,
      weight: float = 10.0,
      eps: float = 1e-8,
      name_scope: str = "gradient_penalty",
  ) -> tf.Tensor:
      """Gradient penalty loss.
  
      This function implements the gradient penalty loss from "Improved Training of
      Wasserstein GANs"
      https://arxiv.org/abs/1704.00028
      This version also implements the multi-scale extension proposed in "MSG-GAN:
      Multi-Scale Gradient GAN for Stable Image Synthesis"
      https://arxiv.org/abs/1903.06048
  
      Args:
        real_data: Samples from the real data.
        generated_data: Samples from the generated data.
        discriminator: The Keras model of the discriminator. This model is expected
          to take as input a tf.Tensor or a sequence of tf.Tensor depending on what
          is provided as generated_data and real_data.
        weight: The weight of the loss.
        eps: A small positive value that is added to the argument of the square root
          to avoid undefined gradients.
        name_scope: The name scope of the loss.
  
      Returns:
        The gradient penalty loss.
  
      Raises:
        TypeError if real_data and generated_data are not both either tf.Tensor or
        both a sequence of tf.Tensor.
        ValueError if the numnber of elements in real_data and generated_data are
        not equal.
      """
      with tf.name_scope(name=name_scope):
          with tf.GradientTape() as tape:
              if isinstance(real_data, tf.Tensor) and isinstance(
                  generated_data, tf.Tensor
              ):
                  epsilon = tf.random.uniform(
                      [tf.shape(real_data)[0]]
                      + [1] * (real_data.shape.ndims - 1),
                      minval=0.0,
                      maxval=1.0,
                      dtype=real_data.dtype,
                  )
                  interpolated_data = (
                      epsilon * real_data
                      + (1.0 - epsilon) * generated_data
                  )
              elif isinstance(real_data, abc.Sequence) and isinstance(
                  generated_data, abc.Sequence
              ):
                  if len(real_data) != len(generated_data):
                      raise ValueError(
                          "The number of elements in real_data and generated_data are "
                          "expected to be equal but got: %d and %d"
                          % (len(real_data), len(generated_data))
                      )
                  epsilon = tf.random.uniform(
                      [tf.shape(real_data[0])[0]]
                      + [1] * (real_data[0].shape.ndims - 1),
                      minval=0.0,
                      maxval=1.0,
                      dtype=real_data[0].dtype,
                  )
                  interpolated_data = [
                      epsilon * real_level
                      + (1.0 - epsilon) * generated_level
                      for real_level, generated_level in zip(
                          real_data, generated_data
                      )
                  ]
              else:
                  raise TypeError(
                      "real_data and generated data should either both be a tf.Tensor "
                      "or both a sequence of tf.Tensor but got: %s and %s"
                      % (type(real_data), type(generated_data))
                  )
              # By default the gradient tape only watches trainable variables.
              tape.watch(interpolated_data)
              interpolated_labels = discriminator(interpolated_data)
  
              with tf.name_scope(name="gradients"):
                  gradients = tape.gradient(
                      target=interpolated_labels,
                      sources=interpolated_data,
                  )
  
          if isinstance(real_data, tf.Tensor):
              gradient_squares = tf.reduce_sum(
                  input_tensor=tf.square(gradients),
                  axis=tuple(range(1, gradients.shape.ndims)),
              )
  
              gradient_norms = tf.sqrt(gradient_squares + eps)
              penalties_squared = tf.square(gradient_norms - 1.0)
              return weight * penalties_squared
          else:
              all_penalties_squared = []
              for gradients_level in gradients:
                  gradient_squares_level = tf.reduce_sum(
                      input_tensor=tf.square(gradients_level),
                      axis=tuple(range(1, gradients_level.shape.ndims)),
                  )
  
                  gradient_norms_level = tf.sqrt(
                      gradient_squares_level + eps
                  )
                  all_penalties_squared.append(
                      tf.square(gradient_norms_level - 1.0)
                  )
  
              return (
                  weight
                  * tf.add_n(all_penalties_squared)
                  * 1.0
                  / len(real_data)
              )
  
file_path: graphics/tensorflow_graphics/projects/gan/losses.py
imports: from collections import abc
  from typing import Sequence, Union
  import tensorflow
label_desc: Use image node_id 1 for calling a global function with example usage: gradient_penalty_loss(real_data, generated_data, discriminator, weight, eps, name_scope) and returns: unknown, unknown