def gradient_penalty_without_bn(net, variables, x):
  """Gradient penalty when each sample is independent.

  if in the training mode
    forward_fn = lambda x: net(train=True).apply(variables, x, mutable=True)
    out, vjp_fn, new_variables = jax.vjp(forward_fn, x, has_aux=True)
    gradient = vjp_fn(jnp.ones_like(out))[0]
    return out, gradient, new_variables
  but as we want to real and fake to have same discriminator parameters, so we
  do not want to update its parameters, so we use train mode.

  !!! discrimintor should not have BN or dropout in this case.
  Args:
    net: The network.
    variables: network variables.
    x: input image.

  Returns:
    out: output logit.
    gradient: The gradient.
  """
  forward_fn = lambda x: net(train=False).apply(variables, x, mutable=False)
  out, vjp_fn = jax.vjp(forward_fn, x, has_aux=False)
  gradient = vjp_fn(jnp.ones_like(out))[0]
  return out, gradient
Gradient penalty when each sample is independent.

  if in the training mode
    forward_fn = lambda x: net(train=True).apply(variables, x, mutable=True)
    out, vjp_fn, new_variables = jax.vjp(forward_fn, x, has_aux=True)
    gradient = vjp_fn(jnp.ones_like(out))[0]
    return out, gradient, new_variables
  but as we want to real and fake to have same discriminator parameters, so we
  do not want to update its parameters, so we use train mode.

  !!! discrimintor should not have BN or dropout in this case.
  Args:
    net: The network.
    variables: network variables.
    x: input image.

  Returns:
    out: output logit.
    gradient: The gradient.
  
videogvt/train_lib/losses.py