text: import tensorflow as tf
  from tensorflow import keras
  
  
  class ConditionalGAN(keras.Model):
      """A light weight Conditional GAN trainer class, just plug in the generator
      and discriminator models and provide the required parameters.
  
      Args:
          discriminator (keras.models.Model): discriminator model with input
              channels = number of channels of the image + number of classes.
          generator (keras.models.Model): discriminator model with input
              dimension = latent_dimension + number of classes.
          latent_dim (int]): latent dimension.
      """
  
      def __init__(
          self,
          discriminator,
          generator,
          latent_dim,
      ):
          super().__init__()
          self.discriminator = discriminator
          self.generator = generator
          self.latent_dim = latent_dim
  
      def compile(self, d_optimizer, g_optimizer, loss_fn):
          """Our very familiar compile function.
  
          Args:
              d_optimizer (keras.optimizers): optimizer function for discriminator.
              g_optimizer (keras.optimizers): optimizer function for generator.
              loss_fn (keras.losses): loss fuction for the GAN.
          """
          super().compile()
          self.d_optimizer = d_optimizer
          self.g_optimizer = g_optimizer
          self.loss_fn = loss_fn
  
      def train_step(self, data):
          # Unpack the data.
          real_images, one_hot_labels = data
          image_size = real_images.shape[1]
          num_classes = one_hot_labels.shape[-1]
  
          # Add dummy dimensions to the labels so that they can be concatenated with
          # the images. This is for the discriminator.
          image_one_hot_labels = one_hot_labels[:, :, None, None]
          image_one_hot_labels = tf.repeat(
              image_one_hot_labels, repeats=[image_size * image_size]
          )
          image_one_hot_labels = tf.reshape(
              image_one_hot_labels, (-1, image_size, image_size, num_classes)
          )
  
          # Sample random points in the latent space and concatenate the labels.
          # This is for the generator.
          batch_size = tf.shape(real_images)[0]
          random_latent_vectors = tf.random.normal(shape=(batch_size, self.latent_dim))
          random_vector_labels = tf.concat(
              [random_latent_vectors, one_hot_labels], axis=1
          )
  
          # Decode the noise (guided by labels) to fake images.
          generated_images = self.generator(random_vector_labels)
  
          # Combine them with real images. Note that we are concatenating the labels
          # with these images here.
          fake_image_and_labels = tf.concat([generated_images, image_one_hot_labels], -1)
          real_image_and_labels = tf.concat([real_images, image_one_hot_labels], -1)
          combined_images = tf.concat(
              [fake_image_and_labels, real_image_and_labels], axis=0
          )
  
          # Assemble labels discriminating real from fake images.
          labels = tf.concat(
              [tf.ones((batch_size, 1)), tf.zeros((batch_size, 1))], axis=0
          )
  
          # Train the discriminator.
          with tf.GradientTape() as tape:
              predictions = self.discriminator(combined_images)
              d_loss = self.loss_fn(labels, predictions)
          grads = tape.gradient(d_loss, self.discriminator.trainable_weights)
          self.d_optimizer.apply_gradients(
              zip(grads, self.discriminator.trainable_weights)
          )
  
          # Sample random points in the latent space.
          random_latent_vectors = tf.random.normal(shape=(batch_size, self.latent_dim))
          random_vector_labels = tf.concat(
              [random_latent_vectors, one_hot_labels], axis=1
          )
  
          # Assemble labels that say "all real images".
          misleading_labels = tf.zeros((batch_size, 1))
  
          # Train the generator (note that we should *not* update the weights
          # of the discriminator)!
          with tf.GradientTape() as tape:
              fake_images = self.generator(random_vector_labels)
              fake_image_and_labels = tf.concat([fake_images, image_one_hot_labels], -1)
              predictions = self.discriminator(fake_image_and_labels)
              g_loss = self.loss_fn(misleading_labels, predictions)
          grads = tape.gradient(g_loss, self.generator.trainable_weights)
          self.g_optimizer.apply_gradients(zip(grads, self.generator.trainable_weights))
  
          return {"d_loss": d_loss, "g_loss": g_loss}
  