code: def gradient_penalty(
      x_real, x_fake, applied_discriminator_fn, rngkey
  ):
      assert len(x_real) == len(x_fake)
      alpha = jax.random.uniform(rngkey, shape=[len(x_real), 1, 1, 1])
      x_hat = x_real * alpha + x_fake * (1 - alpha)
      grads = jax.grad(lambda x: applied_discriminator_fn(x)[0].mean())(
          x_hat
      )
      norm = jnp.sqrt((grads**2).sum(axis=[1, 2, 3]))
      penalty = (norm - 1) ** 2
      return penalty.mean() * LAMBDA_GP
  
file_path: elegy/examples/need-fixing/WGAN-GP/model.py
imports: import jax
  import jax.numpy
  import optax
  import elegy
label_desc: Use image node_id 1 for calling a global function with example usage: gradient_penalty(x_real, x_fake, applied_discriminator_fn, rngkey) and returns: unknown