How is a 3-d tensor indexed by two 2d tensors?
Here's a snapshot from line 15-20 in DIM

def random_permute(X):
    X = X.transpose(1, 2)
    b = torch.rand((X.size(0), X.size(1))).cuda()
    idx = b.sort(0)[1]
    adx = torch.range(0, X.size(1) - 1).long()
    X = X[idx, adx[None, :]].transpose(1, 2)

    return X


where X is a tensor of size [64, 64, 128], idx a tensor of size [64, 64], adx  a tensor of size [64].
How does X = X[idx, adx[None, :]] work? How can we use two 2d tensors to index a 3d tensor? What really happens to X after this indexing?

Things will be more clear if we consider a smaller concrete example. Let

x = np.arange(8).reshape(2, 2, 2)
b = np.random.rand(2, 2)
idx = b.argsort(0) # e.g. idx=[[1, 1], [0, 0]]
adx = np.arange(2)[None, :] # [[0, 1]]
y = x[idx, adx] # implicitly expanding 'adx' to [[0, 1], [0, 1]]


In this example, we'll have y as

y[0, 0] = x[idx[0, 0], adx[0, 0]]=x[1, 0]
y[0, 1] = x[idx[0, 1], adx[0, 1]]=x[1, 1]
y[1, 0] = x[idx[1, 0], adx[1, 0]]=x[0, 0]
...


It may be helpful to see how we do the same in tensorflow:

d0, d1, d2 = x.shape.as_list()
b = np.random.rand(d0, d1)
idx = np.argsort(b, 0)
idx = idx.reshape(-1)
adx = np.arange(0, d1)
adx = np.tile(adx, d0)

y = tf.reshape(tf.gather_nd(x, zip(idx, adx)), (d0, d1, d2))

