import torch
from typing import Tuple


def get_pair_embedding(
    pair_info: torch.Tensor, 
    embedding: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    We are working with N pairs, the k'th pair is (a_k,b_k).

        Args:
        -   pair_info: (N,6) array representing (bi,hi,wi, bj,hj,wj)
                where these represent (batch index, row index, column index)
                into a NCHW tensor for paired samples A and B.
        -   embedding: NCHW tensor representing a minibatch of per-pixel embeddings

        Returns:
        -   a_embedding: (N,C) array representing channels at pixel (i,j) 
                of specific minibatch examples
        -   b_embedding: As above.
    """
    bi = pair_info[:,0]
    hi = pair_info[:,1]
    wi = pair_info[:,2]

    bj = pair_info[:,3]
    hj = pair_info[:,4]
    wj = pair_info[:,5]

    a_embedding = embedding[bi,:,hi,wi]
    b_embedding = embedding[bj,:,hj,wj]
    return a_embedding, b_embedding