Multidimensional gather

1

Hi,

Is there any operation available in PyTorch to do something like the one below?

out[x][y] = in[ind[0][x][y]] [ind[1][x][y]] [ind[2][x][y]]

I want to avoid loops here so I prepared an index tensor like this.

So the context is there are batches of input images(in).

In index tensor(ind), the channel of the target image is stored in 0 dim and where to sample pixels from that target image for corresponding output(out) pixel is stored in 1,2 dim.

2

Could you write a (slow) reference code using nested loops to see what the expected output would be, please?

3

Thanks for your reply. Actually, I just solved the problem by simply flattening the in and ind tensors.

in (b,H,W) flattened to in (b x H x W)
ind (b,H,W) flattened to ind (b x H x W)

Then simply by in[ind], I could get out (H x W) and rearranged values.

Thank you for your interest.