import torch


def aggregate_sum(X, adj, self_loop=False, device='cpu', avg_d=None):
    """
    Aggregate each node by the sum of it's neighbours, weighted by the edge between the nodes.
    Parameters
    --------------
        X: tensor(..., N, N, Din)
            Input feature tensor to aggregate.
        adj: tensor(..., N, N)
            Batches of adjacency matrices of graphs.
    Returns
    -------------
        X_sum: tensor(..., N, N, Din)
            Aggregation results
            
    """

    if self_loop:  # add self connections
        N = adj.shape[-1]
        adj = adj + torch.eye(N, device=device).unsqueeze(0)

    X_sum = torch.sum(torch.mul(X, adj.unsqueeze(-1)), dim=-2)
    return X_sum