import torch


def reduce(x, reduction='mean'):
    """This is a very common line in my custom loss functions so I'm providing
    a convenience function for it. I think this function call is also more
    intuitive than the code behind it.

    Parameters
    ----------
    x: torch.Tensor
        The object to reduce.
    reduction: str
        Should be one of ('mean', 'sum', 'none'), though technically some
        other operations (e.g. 'std') are supported.

    Returns
    -------
    torch.Tensor: Scalar if using 'mean' or 'sum', otherwise the same as
    input `x`.

    Examples
    --------
    def squared_error(x, reduction='mean'):
        return reduce(x.pow(2), reduction)
    """
    return x if reduction == 'none' else getattr(torch, reduction)(x)