def train_g_d(
    train_state: TrainState,
    batch: Batch,
    *,
    model_dict: Dict[str, nn.Module],
    config: ml_collections.ConfigDict) -> Tuple[TrainState, metrics.Collection]:
  """Perform a single training step of g and d."""

  dtype = train_utils.get_dtype(config)
  step = train_state.global_step
  new_rng, _ = jax.random.split(train_state.rng)

  generator = model_dict['generator']
  discriminator = model_dict['discriminator']
  mutable = ['batch_stats', 'spectral_norm_stats']

  if config.pretrained_image_model:
    additional_data = train_utils.get_additional_data()

  def loss_fn(params_d, params_g):
    g_variables = {'params': params_g, **train_state.g_model_state}
    d_variables = {'params': params_d, **train_state.d_model_state}
    real_video = batch['inputs']
    (generated_video, result_dict), new_g_variables = generator.apply(
        g_variables, real_video, is_train=True, mutable=mutable)

    if config.vqgan.gradient_penalty == 'none':
      all_videos = jnp.concatenate([real_video, generated_video])
      # TODO(Lijun-Yu): maybe dual discriminators
      logit, new_d_variables = discriminator.apply(
          d_variables, all_videos, mutable=mutable)
      logit = jnp.asarray(logit, dtype)
      real_logit, fake_logit = jnp.split(logit, 2)
      grad_penalty = 0.0
    elif config.vqgan.gradient_penalty == 'r1':
      real_logit, grad_penalty = losses.r1_gradient_penalty(
          lambda train: discriminator,
          d_variables,
          real_video,
          penalty_cost=config.vqgan.grad_penalty_cost)
      fake_logit, new_d_variables = discriminator.apply(
          d_variables, generated_video, mutable=mutable)
    else:
      raise ValueError(
          f'{config.vqgan.gradient_penalty} is not a recognized gradient'
          ' penalty type.'
      )

    real_logit = jnp.asarray(real_logit, dtype)
    fake_logit = jnp.asarray(fake_logit, dtype)
    real_pred = jnp.mean(real_logit)
    fake_pred = jnp.mean(fake_logit)

    if config.get('lecam_weight', 0.) > 0:
      metadata = train_state.metadata
      lecam_loss = losses.lecam_reg(real_pred, fake_pred,
                                    metadata['lecam_ema_real'],
                                    metadata['lecam_ema_fake'])
      lecam_loss = lecam_loss * config.lecam_weight
    else:
      lecam_loss = 0.0

    d_adversarial_loss = losses.discriminator_loss(
        real_logit=real_logit,
        fake_logit=fake_logit,
        loss_type=config.vqgan.loss_type)
    g_adversarial_loss = losses.generator_loss(
        fake_logit=fake_logit, loss_type=config.vqgan.loss_type
    ) * config.vqgan.g_adversarial_loss_weight

    reconstruction_loss = losses.l2_loss(real_video, generated_video)
    perceptual_loss = 0.0

    if config.perceptual_loss_weight != 0:
      assert config.pretrained_image_model
      # TODO(Lijun-Yu, roadjiang): maybe add a pretrained video model.
      perceptual_loss = losses.calculate_perceptual_loss_on_pretrained(
          additional_data['image_model'],
          additional_data['image_model_state'],
          train_utils.flatten_t_dim(real_video),
          train_utils.flatten_t_dim(generated_video),
          perceptual_loss_on_logit=config.perceptual_loss_on_logit
      ) * config.perceptual_loss_weight

    quantizer_loss = result_dict['quantizer_loss']
    if config.vqvae.get('entropy_loss_enlarge_steps', 0) > 0:
      quantizer_loss += (
          result_dict['entropy_loss']
          * config.vqvae.entropy_loss_enlarge_ratio
          * jnp.maximum(0, 1 - step / config.vqvae.entropy_loss_enlarge_steps)
      )
    logit_laplace_loss = 0.
    d_loss = d_adversarial_loss + grad_penalty + lecam_loss
    g_loss = (
        reconstruction_loss
        + g_adversarial_loss
        + perceptual_loss
        + quantizer_loss
        + logit_laplace_loss
    )
    new_g_state = flax.core.freeze(new_g_variables)
    new_d_state = flax.core.freeze(new_d_variables)
    return (d_loss,
            g_loss), (new_g_state, new_d_state, d_adversarial_loss,
                      grad_penalty, reconstruction_loss, g_adversarial_loss,
                      perceptual_loss, quantizer_loss, logit_laplace_loss,
                      lecam_loss, (real_pred, fake_pred))

  params_d = train_state.d_params
  params_g = train_state.g_params
  (d_loss,
   g_loss), func_vjp, (new_g_state, new_d_state, d_adversarial_loss,
                       grad_penalty, reconstruction_loss, g_adversarial_loss,
                       perceptual_loss, quantizer_loss, logit_laplace_loss,
                       lecam_loss, preds) = jax.vjp(
                           loss_fn, params_d, params_g, has_aux=True)

  d_grad, _ = func_vjp((1., 0.))
  _, g_grad = func_vjp((0., 1.))

  # Compute average gradient across multiple workers.
  d_grad = jax.lax.pmean(d_grad, axis_name='device')
  g_grad = jax.lax.pmean(g_grad, axis_name='device')

  d_updates, new_d_opt_state = train_state.d_tx.update(d_grad,
                                                       train_state.d_opt_state,
                                                       train_state.d_params)
  new_d_params = optax.apply_updates(train_state.d_params, d_updates)

  g_updates, new_g_opt_state = train_state.g_tx.update(g_grad,
                                                       train_state.g_opt_state,
                                                       train_state.g_params)
  new_g_params = optax.apply_updates(train_state.g_params, g_updates)

  ema_decay = config.polyak_decay

  new_ema_params = vqgan_train_lib.compute_ema_params(
      train_state.ema_params, new_g_params, config
  )

  if config.get('lecam_weight', 0.) > 0:
    metadata = train_state.metadata
    real_pred, fake_pred = preds
    new_metadata = {
        'lecam_ema_real':
            metadata['lecam_ema_real'] * ema_decay +
            (1 - ema_decay) * real_pred,
        'lecam_ema_fake':
            metadata['lecam_ema_fake'] * ema_decay +
            (1 - ema_decay) * fake_pred,
    }
  else:
    new_metadata = train_state.metadata

  new_state = train_state.replace(  # pytype: disable=attribute-error
      global_step=step + 1,
      g_opt_state=new_g_opt_state,
      d_opt_state=new_d_opt_state,
      g_params=new_g_params,
      d_params=new_d_params,
      g_model_state=new_g_state,
      d_model_state=new_d_state,
      ema_params=new_ema_params,
      metadata=new_metadata,
      rng=new_rng)
  metrics_update = TrainMetrics.gather_from_model_output(
      axis_name='device',
      g_loss=g_loss,
      d_loss=d_loss,
      d_adversarial_loss=d_adversarial_loss,
      grad_penalty=grad_penalty,
      reconstruction_loss=reconstruction_loss,
      g_adversarial_loss=g_adversarial_loss,
      perceptual_loss=perceptual_loss,
      quantizer_loss=quantizer_loss,
      logit_laplace_loss=logit_laplace_loss,
      lecam_loss=lecam_loss,
  )
  return new_state, metrics_update
Perform a single training step of g and d.
videogvt/trainers/vqgan_trainer.py