index_select_test.py

import numpy as np
import torch
from torch.autograd import Variable

def set_seed(seed):
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)

def main():
    set_seed(1)
    cuda = True
    n, e, c = 5, 100, 5
    input = torch.randn(n,c)
    idxn = torch.from_numpy(np.random.randint(n,size=e))  # indices are repeated

    if cuda:
        input = input.cuda(); idxn = idxn.cuda()

    gradsI, gradsS = [], []
    N = 2
    for i in range(N):
        inputv = Variable(input, requires_grad=True)
        sel_input = torch.index_select(inputv, 0, Variable(idxn))
        sel_input.retain_grad()

        # the following computation is one of the "random conditions"
        data = [torch.sum(sel_input.narrow(0,0,e//2), 0),
                torch.sum(sel_input.narrow(0,e//2,e//2), 0) ]
        out = torch.cat(data,0)
        out.exp().sum().backward()

        gradsI.append(inputv.grad.data.cpu().clone())
        gradsS.append(sel_input.grad.data.cpu().clone())

    for i in range(N):
        for j in range(N):
            print('-------------------')
            print(i,j,(gradsI[i]-gradsI[j]).abs().max(), (gradsS[i]-gradsS[j]).abs().max())
            print(gradsI[i][0,2], gradsI[j][0,2])

if __name__ == '__main__':
    main()