How to index a list with a TensorFlow tensor?

Question

Assume a list with non concatenable objects which needs to be accessed via a look up table. So the list index will be a tensor object but this is not possible.

 tf_look_up = tf.constant(np.array([3, 2, 1, 0, 4]))
 index = tf.constant(2)
 list = [0,1,2,3,4]

 target = list[tf_look_up[index]]

This will bring out the following error message.

 TypeError: list indices must be integers or slices, not Tensor

Is the a way/workaround to index lists with tensors?

@YaroslavBulatov what if the list is produced dynamically, e.g. the states that a RNN produces. Any way list dynamic indexing would work? — chentingpc, May 16 '17 at 16:01

score 14 · Answer 1 · answered Apr 28 '17 at 06:37

14

tf.gather is designed for this purpose.

Simply run tf.gather(list, tf_look_up[index]), you'll get what you want.

answered Apr 28 '17 at 06:37

soloice

980
8
17

1

What if the list is something like `[[1], [2, 3], [4]]`? – Ben Nov 23 '21 at 19:29

score 2 · Answer 2 · answered Mar 23 '17 at 13:22

Tensorflow actually has support for a HashTable. See the documentation for more details.

Here, what you could do is the following:

table = tf.contrib.lookup.HashTable(
    tf.contrib.lookup.KeyValueTensorInitializer(tf_look_up, list), -1)

Then just get the desired input by running

target = table.lookup(index)

Note that -1 is the default value if the key is not found. You may have to add key_dtype and value_dtype to the constructor depending on the configuration of your tensors.

How to index a list with a TensorFlow tensor?

2 Answers2

Linked