Indexing a Tensor only partially with LongTensors

1

Hi there,
I’m trying something along those lines:
I have indices to a 3D volume (b), and b.nonzero() gives me (N, 3) indices to index with.
Now I have another 3D volume with some more dimensions (a) and I would like to get something of the shape (BS, F, N) with N being the number of nonzero samples in b.

a = torch.randn(1,2, 10,10,10)
b = torch.randn(     10,10,10)
a[:, :, (b > 0).nonzero(as_tuple=True)]

Leading to

TypeError: only integer tensors of a single element can be converted to an index

Obviously b[(b > 0).nonzero(as_tuple=True)] works, but I can’t seem to index the first two dimensions differently (i.e. with something else but a LongTensor).

The only way I have in mind to do this would be a torch.nn.functional.grid_sample() with nearest mode after dividing the indices by the respective dimension lengths, but that seems unnecessary and slow.

Any suggestions on how I can achieve the above behavior?

2

Hi Dominik!

Use pytorch tensor indexing as follows:

Create five appropriately-structured index tensors – one each for each dimension
of a – of shape [BS, F, N], and use these to index into a.

Thus:

>>> import torch
>>> print (torch.__version__)
1.12.0
>>>
>>> _ = torch.manual_seed (2022)
>>>
>>> a = torch.randn (1, 2, 10, 10, 10)
>>> b = torch.randn (10, 10, 10)
>>>
>>> BS = a.size (0)
>>> F = a.size (1)
>>> N = (b > 0).sum()
>>>
>>> ind_n0, ind_n1, ind_n2 = (b > 0).nonzero (as_tuple = True)
>>>
>>> ind_bs = torch.arange (BS)[:, None, None].expand (BS, F, N)
>>> ind_f = torch.arange (F)[None, :, None].expand (BS, F, N)
>>> ind_n0 = ind_n0[None, None, :].expand (BS, F, N)
>>> ind_n1 = ind_n1[None, None, :].expand (BS, F, N)
>>> ind_n2 = ind_n2[None, None, :].expand (BS, F, N)
>>>
>>> result = a[ind_bs, ind_f, ind_n0, ind_n1, ind_n2]
>>>
>>> result.shape
torch.Size([1, 2, 511])
>>>
>>> # check result
>>> torch.equal (result[0, 0], a[0, 0][(b > 0).nonzero (as_tuple = True)])
True
>>> torch.equal (result[0, 1], a[0, 1][(b > 0).nonzero (as_tuple = True)])
True

Best.

K. Frank

1 Like