import torch


def mean_tensor(input, norm_ax, keepdim=True):
    """
    outputs the mean of a pytorch tensor object along axis
    """
    assert isinstance(norm_ax, (tuple, int)), f'norm_ax must be a tuple or int, {norm_ax} is a {type(norm_ax)}'
    dims = input.size()

    if isinstance(norm_ax, int):
        n = dims[norm_ax]
    else:
        n = 1
        for d in norm_ax: n *= dims[d]

    mu = torch.sum(input, dim=norm_ax, keepdim=keepdim) / n

    return mu