import torch


def kinetic_energy_reg_func(x: torch.Tensor,
                            t: torch.Tensor,
                            dx: torch.Tensor,
                            unused_context) -> torch.Tensor:
    """
    Quadratic cost / kinetic energy regularization: https://arxiv.org/pdf/2002.02798.pdf

    :param x:
    :param t:
    :param dx:
    :return: The regularization term.
    """

    del x, t
    dx = dx.view(dx.shape[0], -1)
    return 0.5 * dx.pow(2).mean(dim=-1)