import torch


def av_norm(tensor, average_logs=False):
    """
    Return the average of the norms (or log norms) of a batch of tensors,
    with batch index assumed to be zeroth index
    """
    num_inds = len(tensor.shape)
    norms = (tensor ** 2).sum(list(range(1, num_inds)))
    if average_logs: norms = torch.log(norms)
    assert len(norms.shape) == 1
    return norms.mean()