import torch


def get_neighbor_bonds(edge_index, bond_type):
    """
    Takes the edge indices and bond type and returns dictionary mapping atom index to neighbor bond types
    Note: this only includes atoms with degree > 1
    """
    start, end = edge_index
    idxs, vals = torch.unique(start, return_counts=True)
    vs = torch.split_with_sizes(bond_type, tuple(vals))
    return {k.item(): v for k, v in zip(idxs, vs) if len(v) > 1}