import torch


def average_aggregator(embedding_batch, embed_dim=0):
    """
    aggregates by averaging embedding lists in each of the n batch entries
    :param embedding_batch: python list of n tensors of m x K-dim embedding vectors (m can vary)
    :param embed_dim: K, used if the embedding_lists is empty (n = 0) to return a (0, K) empty tensor
    :return: n x K tensor of averaged embeddings
    """
    if len(embedding_batch):
        batch_means = [t.mean(dim=0) for t in embedding_batch]
        return torch.stack(batch_means)
    else:
        return torch.tensor([]).view(-1, embed_dim)