How to use a tensor of indices to copy values between two other tensors

Suppose I have the following tensors:

x = 
[[1,2,3],
 [4,5,6],
 [7,8,9]]

y = 
[[11,21,31],
 [41,51,61],
 [71,81,91]]

z=
[[1],
 [2],
 [0]]

Where x,y are both (h,w) and z is (h,1). I want to create a new tensor k that it is a copy of x except I want to copy the z[i][0] column at row i from y into row i at column z[i][0] in x. Thus, using the tensors from above, k should be:

k = 
[[1,21,3],
 [4,5,61],
 [71,8,9]]

I though using k = x.scatter_(1, z, y) would accomplish this, but it seems to only use values from the first column of y, so I obviously don’t know how scatter_ works. How can I leverage PyTorch to accomplish this so I can avoid using a for-loop?

A combination of gather and scatter_ should work:

x = torch.tensor([[1,2,3],
                  [4,5,6],
                  [7,8,9]])

y = torch.tensor([[11,21,31],
                  [41,51,61],
                  [71,81,91]])

z = torch.tensor([[1],
                  [2],
                  [0]])


k = x.scatter_(1, z, y.gather(1, z))
print(k)
> tensor([[ 1, 21,  3],
          [ 4,  5, 61],
          [71,  8,  9]])
1 Like

This is awesome! Didn’t even know about gather.

1 Like