import torch


def form_pair_info_tensor(
    batch_dim_idxs: torch.Tensor, 
    px_1d_y: torch.Tensor, 
    px_1d_x: torch.Tensor
    ):
    """ Hstack 3 length-N 1d arrays into a (N,3) array

        Args:
        -   batch_dim_idxs: size (N,) array representing indices
                of examples in a minibatch
        -   px_1d_y: size (N,) array representing row indices
        -   px_1d_x: size (N,) array representing column indices

        Returns:
        -   pair_info: (N,3) array
    """
    # batch dim
    N = batch_dim_idxs.shape[0]
    assert batch_dim_idxs.shape == (N,)
    assert px_1d_y.shape == (N,)
    assert px_1d_x.shape == (N,)

    pair_info = torch.stack([batch_dim_idxs, px_1d_y, px_1d_x])
    return pair_info.t()