import torch


def _gradient_penalty(x, descriminator, layers, alpha):
    """
    Compute the gradient penalty term used in WGAN-gp.
    Returns the gradient penalty for each batch entry, the loss term is computed as the average.

    Works by sampling points on the line segment between x_r and x_g, then computing the gradient
    of the critic with respect to each sample point.
    """
    batch_size = x.size(0)
    x.requires_grad = True

    # We compute the gradient of the parameters using the regular autograd.
    # The key to making this work is including `create_graph`, this means that the computations
    # in this penalty will be added to the computation graph for the loss function, so that the
    # second partial derivatives will be correctly computed.
    f_x = descriminator(x, layers, alpha)

    grad = torch.autograd.grad(
        outputs=f_x,
        inputs=x,
        grad_outputs=torch.ones_like(f_x),
        create_graph=True,
        only_inputs=True,
    )[0]

    grad_x_flat = grad.view(batch_size, -1)
    gradient_norm = torch.linalg.norm(grad_x_flat, dim=1)
    gp = torch.pow((gradient_norm - 1.0), 2)

    return gp, grad