Pytorch: How to access tensor(values) by tensor(keys) in python dictionary

Question

I have a dictionary with tensor keys and tensor values. I want to access the values by the keys.

from torch import tensor
x = {tensor(0): [tensor(1)], tensor(1): [tensor(0)]}
for i in x.keys():
  print(i, x[i]) 

Returns:

tensor(0) [tensor(1)]
tensor(1) [tensor(0)]

But when i try to access the values without looping through the keys,

try:
    print(x[tensor(0)])

except:
    print(Exception)
    print(x[0])

Throws Exception:

 KeyError                                   Traceback (most recent call last)
 <ipython-input-34-746d28dcd450> in <module>()
  6 try:
  ----> 7   print(x[tensor(0)])
  8 

  KeyError: tensor(0)

  During handling of the above exception, another exception occurred:

  KeyError                                  Traceback (most recent call last)
  <ipython-input-34-746d28dcd450> in <module>()
  9 except:
  10   print(Exception)
  ---> 11   print(x[0])
  12 continue

  KeyError: 0

Answer 1

In PyTorch, hashes of Tensors are a function of their id, not the actual value. Because Python dictionaries use the hashes for lookup, lookup fails. See this Github discussion.

In [4]: hash(tensor(0)) == hash(tensor(0))                                      
Out[4]: False

In [5]: hash(tensor(0))                                                         
Out[5]: 4364730928

In [6]: hash(tensor(0))                                                         
Out[6]: 4362187312

In [7]: hash(tensor(0))                                                         
Out[7]: 4364733808

In order to achieve what you want, you could either use plain Python integers as keys, or use an Embedding object as x.

Answer 2

There's at least three issues here.

1. If x is a dictionary of tensor keys, then of course x[0] will not work. 0 is not a key of it. Hence the inner KeyError that occured during the other exception.
2. Not actually relevant to your errors, but print(Exception) is almost surely not what you want. It prints the class object (if that is the right term) of the class Exception. You likely rather meant

   except Exception as e:
     print(e)
   

   or more specifically, except KeyError (otherwise it will just catch every kind of exception).

3. The real thing: you don't want to use a tensor as key in the first place. It's a mutable type, compared by reference, not by value. Every tensor(something) call will create a new object, hashing to a different value than the tensor(something) you specified as key.

   Use the actual integers instead.