Pytorch: How can I find indices of first nonzero element in each row of a 2D tensor?

Question

I have a 2D tensor with some nonzero element in each row like this:

import torch
tmp = torch.tensor([[0, 0, 1, 0, 1, 0, 0],
                    [0, 0, 0, 1, 1, 0, 0]], dtype=torch.float)

I want a tensor containing the index of first nonzero element in each row:

indices = tensor([2],
                 [3])

How can I calculate it in Pytorch?

score 11 · Answer 1 · edited Jul 26 '20 at 08:43

11

I have simplified Iman's approach to do the following:

idx = torch.arange(tmp.shape[1], 0, -1)
tmp2= tmp * idx
indices = torch.argmax(tmp2, 1, keepdim=True)

edited Jul 26 '20 at 08:43

Ash

4,353
5
27
37

answered Feb 13 '20 at 07:46

Shay

111
1
2

this is such an easy but smart solution! – Tian Jul 03 '20 at 23:54
note that if there are no non-zeros this defaults to zero as the first non-zero index. – sachinruk Dec 19 '21 at 07:44

Iman Aliabdi · Accepted Answer · 2019-05-12T14:07:23.413

I could find a tricky answer for my question:

  tmp = torch.tensor([[0, 0, 1, 0, 1, 0, 0],
                     [0, 0, 0, 1, 1, 0, 0]], dtype=torch.float)
  idx = reversed(torch.Tensor(range(1,8)))
  print(idx)

  tmp2= torch.einsum("ab,b->ab", (tmp, idx))

  print(tmp2)

  indices = torch.argmax(tmp2, 1, keepdim=True)
  print(indeces)

The result is:

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

score 1 · Answer 3 · answered Nov 15 '21 at 07:59

1

All the nonzero values are equal, so argmax returns the first index.

tmp = torch.tensor([[0, 0, 1, 0, 1, 0, 0],
                    [0, 0, 0, 1, 1, 0, 0]])
indices = tmp.argmax(1)

answered Nov 15 '21 at 07:59

Seppo Enarvi

2,860
3
28
25

Pytorch: How can I find indices of first nonzero element in each row of a 2D tensor?

3 Answers3

Linked