Using .index_put() with a 3D tensor

1

if I have a 2D tensor like this one:

>>> torch.ones((2, 4))
[[1, 1, 1, 1],
 [1, 1, 1, 1]]

and want to fill two positions per row with 0, to get:

[[1, 0, 1, 0],
 [0, 1, 1, 0]]

I can do:

torch.ones((2, 4)).index_put((torch.arange(2).unsqueeze(1), torch.LongTensor([[1,3], [0,3]])), torch.Tensor([0]))

What about a 3D tensor? Let’s say I want to fill in a torch.ones(2, 3, 4) tensor with some zeros, to get:

tensor([[[1., 0., 1., 0.],
         [0., 1., 1., 0.],
         [1., 0., 0., 1.]],

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

if I have the zero-indices stored as:

torch.LongTensor([[[1,3],
                   [0,3],
                   [1,2]],

                  [[0,2],
                   [1,2],
                   [2,3]]])

is there a way to use these indices, to tell .index_put() where to place the zeros?

2

Hi,

For any number of dimensions, you can use scatter to achieve this:

import torch

ind= torch.tensor([[[1,3],
                   [0,3],
                   [1,2]],

                  [[0,2],
                   [1,2],
                   [2,3]]])

base = torch.ones(2, 3, 4)

base.scatter_(2, ind, 0)

print(base)
3

Hi Alban!

That’s exactly what I needed, thanks a lot

1 Like