def get_pair(tensor):
    """Gets paired combination.

    Given a tensor of shape (x0, x1, x2), this gives you a tensor of shape
    (x0, x1, x1, x2) where the pairs of elements in the x1 dim.
    """
    return tensor.unsqueeze(1) + tensor.unsqueeze(-2)