import torch


def attention_aggregator(embedding_lists, weights, embed_dim=0) -> torch.Tensor:
        """
        Returns a weighted sum of embeddings
        :param embedding_lists: list of n tensors of shape (l, K) embedding tensors (l can vary)
        :param weights: list of n tensors of shape (l,) weights (l can vary, but matches embeddings)
        :param embed_dim: K, used if the embedding_lists is empty (n = 0) to return a (0, K) empty tensor
        :return: weighted sum of embeddings, shape (n, K)
        """
        assert len(embedding_lists) == len(weights), f"aggregation weights different length to embeddings! weights len " \
                                                     f"{len(weights)}, embeds len {len(embedding_lists)}"
        if len(embedding_lists):
            if len(embedding_lists) == 1:
                return embedding_lists[0]  # (1, K) tensor with the single embedding (ie: no aggregation necessary)
            aggregated = torch.stack([torch.sum(emb * w.view(-1, 1), dim=0) for emb, w in zip(embedding_lists, weights)])
            return aggregated  # (n, K) tensor of aggregated embeddings
        else:
            return torch.tensor([]).view(-1, embed_dim)