Scatter with two index tensor

vision
1

Dear all,

suppose I have two index Tensor

index_1, size is [N], and index_2 size is also [N].

Now I have a target tensor, and src tensor like [N, C], what I want is

target[index_1[i], index_2[i], j] = src[i, j]

i is in range N and j is in range C

The most similar function I found is torch.scatter(),

Tensor. scatter_ (dim , index , src , reduce=None ) works like

target[index[i, j], j] = src[i, j]

but it supports only one index tensor. Do you have any alternatives? Because N is quite large, looping is not realistic. I have also tried by indexing

target[index_1, index_2, :] = src

It not raise error but just not correct

Thanks

2

Could you post a minimal code snippet showing the desired result using your “slow” approach?
I’m unsure if this would work:

N, C = 2, 3

target = torch.zeros(N, N, C)
src = torch.randn(N, C)
index1 = torch.randint(0, N, (N,))
index2 = torch.randint(0, N, (N,))

target[index1.unsqueeze(1), index2] = src

so I would like to compare it to your reference.