def lecam_reg(real_pred, fake_pred, ema_real_pred, ema_fake_pred):
  """Lecam loss for data-efficient and stable GAN training.

  Described in https://arxiv.org/abs/2104.03310

  Args:
    real_pred: Prediction (scalar) for the real samples.
    fake_pred: Prediction for the fake samples.
    ema_real_pred: EMA prediction (scalar)  for the real samples.
    ema_fake_pred: EMA prediction for the fake samples.

  Returns:
    Lecam regularization loss (scalar).
  """
  assert real_pred.ndim == 0 and ema_fake_pred.ndim == 0
  lecam_loss = jnp.mean(jnp.power(jax.nn.relu(real_pred - ema_fake_pred), 2))
  lecam_loss += jnp.mean(jnp.power(jax.nn.relu(ema_real_pred - fake_pred), 2))
  return lecam_loss
Lecam loss for data-efficient and stable GAN training.

  Described in https://arxiv.org/abs/2104.03310

  Args:
    real_pred: Prediction (scalar) for the real samples.
    fake_pred: Prediction for the fake samples.
    ema_real_pred: EMA prediction (scalar)  for the real samples.
    ema_fake_pred: EMA prediction for the fake samples.

  Returns:
    Lecam regularization loss (scalar).
  
videogvt/train_lib/losses.py