import torch


def reduce_tensor(vec: torch.Tensor, reduction: str = 'mean'):
    """Global reduction of tensor based on str

    Args:
        vec: torch.FloatTensor
        reduction: str, one of ['sum', 'mean', 'none'], default 'mean'
    """
    assert reduction in ['sum', 'mean', 'none']
    if reduction == 'mean':
        return vec.mean()
    elif reduction == 'sum':
        return vec.sum()
    elif reduction == 'none':
        return vec