import torch


def le(a, b, square=False, **kwargs):
    """
    Encodes the loss function for "a <= b". If square is false d = |a - b| is
    used, else d = (a - b)^2.
    """
    if square:
        return torch.clamp((a - b).sign() * (a - b) * (a - b), min=0)
    else:
        return torch.clamp(a - b, min=0)