import torch


def earth_mover_distance(input: torch.Tensor, target: torch.Tensor, r: float = 2):
    """
    Batch Earth Mover's Distance implementation.
    Args:
        input: B x num_classes
        target: B x num_classes
        r: float, to penalize the Euclidean distance between the CDFs

    Returns:

    """
    N, num_classes = input.size()
    input_cumsum = torch.cumsum(input, dim=-1)
    target_cumsum = torch.cumsum(target, dim=-1)

    diff = torch.abs(input_cumsum - target_cumsum) ** r

    class_wise = (torch.sum(diff, dim=-1) / num_classes) ** (1. / r)
    scalar_ret = torch.sum(class_wise) / N
    return scalar_ret