Set 2D tensor elements based on 2D index tensor

vision
1

I’m trying to rewrite the answer to this stack thread in native PyTorch language (as I don’t want to be flicking back and forth between numpy and PyTorch objects.

However I’m struggling to make it work in tensor form, so I was wondering if anyone could give me a hand in seeing where I’m going wrong?

The new form of the question in the thread goes as follows: Say you have the arrays as described in the thread as tensors…

z = torch.tensor([[[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]],

       [[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]],

       [[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]]])


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


v = torch.tensor([[5., 5., 0., 4.],
       [4., 6., 8., 3.],
       [4., 0., 4., 8.],
       [7., 6., 5., 7.]])

And you wish to perform the answer operation:

z = np.array([*map(i.__eq__,np.unique(i))])*np.array([v]*z.shape[0])

…but in a PyTorch way, how does one alter this?

I have tried:

z = torch.tensor([*map(it.__eq__,torch.unique(it))])*torch.stack([vt]*zt.shape[0])

…but I still get an error relating to the use of a list inside torch.tensor() procedure I believe. You can see I’ve altered the second multiplication term to use torch.stack() instead of a list multiplication but I can’t see how to fix the list error.

Error:

ValueError: only one element tensors can be converted to Python scalars

Any help on this matter would be greatly appreciated, thanks in advance!

3

Solved with:

z = torch.stack([*map(i.__eq__,torch.unique(i))]).type(torch.float)*(torch.stack([v]*z.shape[0])).type(torch.float)
4

Hi,

I think you just want to use scatter for this :slight_smile:

i.unsqueeze_(0)
v.unsqueeze_(0)
z.scatter_(0, i, v)