Index into a pytorch tensor and return array of indices of rows

1

Sorry for the basic question, but I’m hoping the nice people here can help me:
Given the below tensor that has vectors of all zeros and vectors with ones and zeros:

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

How can I have an array of indices of the vectors/rows with ones and zeros so the output is like this:

indices = tensor([ 1, 3, 5, 6,...])
2

It looks like there is an implicit reduction here; are you treating rows with multiple ones the same as those that have a single one?
In that case, you could try:

import torch
a = torch.tensor([[0., 0., 0., 0.],
    [0., 1., 1., 0.],
    [0., 0., 0., 0.],
    [0., 0., 1., 0.],
    [0., 0., 0., 0.],
    [0., 0., 1., 0.],
    [1., 0., 0., 1.],
    [0., 0., 0., 0.]])
print(a)
s = torch.sum(a, dim=-1)
print(s)
inds = torch.nonzero(s).reshape(-1)
print(inds)

tensor([[0., 0., 0., 0.],
        [0., 1., 1., 0.],
        [0., 0., 0., 0.],
        [0., 0., 1., 0.],
        [0., 0., 0., 0.],
        [0., 0., 1., 0.],
        [1., 0., 0., 1.],
        [0., 0., 0., 0.]])
tensor([0., 2., 0., 1., 0., 1., 2., 0.])
tensor([1, 3, 5, 6])
1 Like