import torch


def add_self_loops(edge_index, edge_attr, weight=1.0):
    """
    Add a self-loop for each vertex of the graph.

    Args:
        edge_index (`torch.LongTensor`): indices of graph's edges.
        edge_attr (`torch.FloatTensor`): attributes of graph's edges.
        weight (float, optional): weight of a self-loop. Defaults to 1.0.

    Returns:
        (`torch.LongTensor`): indices of graph's edges.
        (`torch.FloatTensor`): attributes of graph's edges.
    """

    self_loop_index = edge_index[0].unique().unsqueeze(0).repeat(2, 1)
    self_loop_attr = weight * torch.ones(self_loop_index.shape[1])

    edge_index = torch.cat((self_loop_index, edge_index), dim=1)
    edge_attr = torch.cat((self_loop_attr, edge_attr))

    return edge_index, edge_attr