How Pytorch Tensor get the index of specific value

Question

With python lists, we can do:

a = [1, 2, 3]
assert a.index(2) == 1

How can a pytorch tensor find the .index() directly?

This has been subsequently answered here: stackoverflow.com/a/51704350/799988 — Matthew
– Matthew, Commented Oct 9, 2019 at 15:11

Manuel Lagunas · Accepted Answer · 2021-02-10 10:54:38Z

100

I think there is no direct translation from list.index() to a pytorch function. However, you can achieve similar results using tensor==number and then the nonzero() function. For example:

t = torch.Tensor([1, 2, 3])
print ((t == 2).nonzero(as_tuple=True)[0])

This piece of code returns

1

[torch.LongTensor of size 1x1]

edited Feb 10, 2021 at 10:54

answered Dec 18, 2017 at 11:26

Manuel Lagunas

2,7911 gold badge22 silver badges32 bronze badges

Sign up to request clarification or add additional context in comments.

4 Comments

Charlie Parker Over a year ago

what would happen if there is no match? how would you handle that?

secsilm Over a year ago

@CharlieParker If there is no match, it will return empty tensor tensor([], dtype=torch.int64)

Karan Shah Over a year ago

How can we extend this to get batch indices? In this case, I'd want indices of values torch.Tensor([1, 2, 3]) all at once, not just 2. Is there a way without for-loops?

user3689574 Over a year ago

@KaranShah, sure, you can use pytorch.org/docs/stable/generated/torch.eq.html for multidirectional elementwise equal.

dopexxx · Accepted Answer · 2021-04-20 09:02:39Z

27

For multidimensional tensors you can do:

(tensor == target_value).nonzero(as_tuple=True)

The resulting tensor will be of shape number_of_matches x tensor_dimension. For example, say tensor is a 3 x 4 tensor (that means the dimension is 2), the result will be a 2D-tensor with the indexes for the matches in the rows.

tensor = torch.Tensor([[1, 2, 2, 7], [3, 1, 2, 4], [3, 1, 9, 4]])
(tensor == 2).nonzero(as_tuple=False)
>>> tensor([[0, 1],
        [0, 2],
        [1, 2]])

answered Apr 20, 2021 at 9:02

dopexxx

2,7061 gold badge24 silver badges30 bronze badges

3 Comments

Charlie Parker Over a year ago

Most complete answer! general beyond just flat tensors.

Charlie Parker Over a year ago

what would happen if there is no match? how would you handle that?

dopexxx Over a year ago

In that case you'd get an empty tensor. The tensor shape would still adhere to the dimension of the input tensor, so in the above example, searching for a 8 would result in an (empty) tensor of shape 0 x 2.

Minions · Accepted Answer · 2022-04-11 19:26:57Z

4

x = torch.Tensor([11, 22, 33, 22])
print((x==22).nonzero().squeeze())

tensor([1, 3])

answered Apr 11, 2022 at 19:26

Minions

5,5376 gold badges56 silver badges105 bronze badges

Comments

Manoj Mohan · Accepted Answer · 2023-09-05 08:59:14Z

3

a = torch.tensor([1, 2, 3])
torch.where(a == 2)[0]
>>>tensor([1])

answered Sep 5, 2023 at 8:59

Manoj Mohan

6,0741 gold badge24 silver badges26 bronze badges

Comments

Frank Xu · Accepted Answer · 2021-05-12 13:45:15Z

2

Based on others' answers:

t = torch.Tensor([1, 2, 3])
print((t==1).nonzero().item())

answered May 12, 2021 at 13:45

Frank Xu

813 bronze badges

1 Comment

Nuclear241 Over a year ago

This is okay if the tensor only contains one occurrence of the intended number. That is because .item() method can only be called on one-element tensor, otherwise it will raise an error.

Charlie Parker · Accepted Answer · 2021-07-16 16:05:00Z

The answers already given are great but they don't handle when I tried it when there is no match. For that see this:

def index(tensor: Tensor, value, ith_match:int =0) -> Tensor:
    """
    Returns generalized index (i.e. location/coordinate) of the first occurence of value
    in Tensor. For flat tensors (i.e. arrays/lists) it returns the indices of the occurrences
    of the value you are looking for. Otherwise, it returns the "index" as a coordinate.
    If there are multiple occurences then you need to choose which one you want with ith_index.
    e.g. ith_index=0 gives first occurence.

    Reference: https://stackoverflow.com/a/67175757/1601580
    :return:
    """
    # bool tensor of where value occurred
    places_where_value_occurs = (tensor == value)
    # get matches as a "coordinate list" where occurence happened
    matches = (tensor == value).nonzero()  # [number_of_matches, tensor_dimension]
    if matches.size(0) == 0:  # no matches
        return -1
    else:
        # get index/coordinate of the occurence you want (e.g. 1st occurence ith_match=0)
        index = matches[ith_match]
        return index

credit to this great answer: https://stackoverflow.com/a/67175757/1601580

jef · Accepted Answer · 2022-07-29 06:35:04Z

1

In my opinion, calling tolist() is simple and easy to understand.

t = torch.Tensor([1, 2, 3])
t.tolist().index(2) # -> 1

answered Jul 29, 2022 at 6:35

jef

4,09311 gold badges46 silver badges80 bronze badges

Comments

vlad · Accepted Answer · 2018-10-14 16:29:20Z

0

Can be done by converting to numpy as follows

import torch
x = torch.range(1,4)
print(x)
===> tensor([ 1.,  2.,  3.,  4.]) 
nx = x.numpy()
np.where(nx == 3)[0][0]
===> 2

edited Oct 14, 2018 at 16:29

answered Jul 14, 2018 at 15:18

vlad

3043 silver badges6 bronze badges

2 Comments

Victor Zuanazzi Over a year ago

Be aware that if you convert to numpy you lose the gradient graph.

Charlie Parker Over a year ago

what would happen if there is no match? how would you handle that?

RANJITH T · Accepted Answer · 2020-08-12 12:10:21Z

-2

for finding index of an element in 1d tensor/array Example

mat=torch.tensor([1,8,5,3])

to find index of 5

five=5

numb_of_col=4
for o in range(numb_of_col):
   if mat[o]==five:
     print(torch.tensor([o]))

To find element index of a 2d/3d tensor covert it into 1d #ie example.view(number of elements)

Example

mat=torch.tensor([[1,2],[4,3])
#to find index of 2

five = 2
mat=mat.view(4)
numb_of_col = 4
for o in range(numb_of_col):
   if mat[o] == five:
     print(torch.tensor([o]))

edited Aug 12, 2020 at 12:10

answered Aug 10, 2020 at 23:01

RANJITH T

11 bronze badge

1 Comment

Simas Joneliunas Over a year ago

This is an old question and the OP is most likely not looking for a quick answer. If you want to post an answer, please take your time to explain what you code does and how does it add to the already existing answers. The code block as it is now is of very little value to the community

Giang Nguyen · Accepted Answer · 2019-05-03 02:03:23Z

-3

For floating point tensors, I use this to get the index of the element in the tensor.

print((torch.abs((torch.max(your_tensor).item()-your_tensor))<0.0001).nonzero())

Here I want to get the index of max_value in the float tensor, you can also put your value like this to get the index of any elements in tensor.

print((torch.abs((YOUR_VALUE-your_tensor))<0.0001).nonzero())

answered May 3, 2019 at 2:03

Giang Nguyen

4981 gold badge8 silver badges17 bronze badges

1 Comment

lericson Over a year ago

FYI torch.isclose(your_value, your_tensor) is a better way to do this.

Mohanraj · Accepted Answer · 2019-03-22 10:25:34Z

-4

    import torch
    x_data = variable(torch.Tensor([[1.0], [2.0], [3.0]]))
    print(x_data.data[0])
    >>tensor([1.])

answered Mar 22, 2019 at 10:25

Mohanraj

1

Collectives™ on Stack Overflow

How Pytorch Tensor get the index of specific value

11 Answers 11

4 Comments

3 Comments

Comments

Comments

1 Comment

Comments

Comments

2 Comments

1 Comment

1 Comment

Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

11 Answers 11

4 Comments

3 Comments

Comments

Comments

1 Comment

Comments

Comments

2 Comments

1 Comment

1 Comment

Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related