How to set a 3-d tensor to 0 according to the value of a 2-d tensor

user_19434393 0
Viewed 28

Suppose I have a 3-d tensor P shaped (B, N, d) and a 2-d tensor Q shaped (B,N), where values in Q are smaller than d. I want to set some values in P to 0 using the indices from Q:

For instance,

P = torch.randn(2,3,6)
Q = torch.tensor([[0,3,4], [2,1,3]])

How can I set P[0,0,0]=0; P[0,1,3]=0; P[0,2,4]=0; P[1,0,2]=0; P[1,1,1]=0; P[1,2,3]=0 using Q and keep other values in P unchanged?

tensor pytorch indexing python
1 Answers

What you are looking to do is:

P[i, j, Q[i, j]] = 0

This is the perfect use case for torch.Tensor.scatter, which has the effect of placing an arbitrary value at designated positions.

We first need the input and indexer to have matching shapes:

>>> Q_ = Q[...,None].expand_as(P)

And apply the scatter function on dim=2, with the hidden value argument...

P.scatter(dim=2, index=Q_, value=0)
tensor([[[0, 2, 5, 4, 3, 7],
         [4, 1, 5, 0, 3, 8],
         [7, 8, 4, 3, 0, 1]],

        [[7, 1, 0, 1, 1, 3],
         [5, 0, 5, 3, 0, 5],
         [6, 3, 5, 0, 7, 2]]])
user_63313690
Related