How to iterate over a tf.tensor in eager mode

Question

I'm trying to iterate over a tensor in eager mode but I can't.

Naturally, you would do something like:

probs = tf.convert_to_tensor(np.array([[1,2,3], [4,5,6], [7,8,9]]))
indexs = tf.convert_to_tensor(np.array([1, 2, 3]))

@tf.function
def iterate_tensor(probs, indexs):
    return [output[label] for output, label in zip(probs, indexs)]
iterate_tensor(probs, indexs)

But this gives the error OperatorNotAllowedInGraphError: iterating over `tf.Tensor` is not allowed:

Another thing that I tried was:

probs = tf.convert_to_tensor(np.array([[1,2,3], [4,5,6], [7,8,9]]))
indexs = tf.convert_to_tensor(np.array([1, 2, 3]))

@tf.function
def iterate_tensor(probs, indexs):
    return tf.map_fn(lambda i: i[0][i[1]], (probs, indexs), dtype=(tf.int64, tf.int64))

iterate_tensor(probs, indexs)

Gives the error ValueError: The two structures don't have the same nested structure.

user_8527630 · Answer

This seems to work:

probs = tf.convert_to_tensor(np.array([[1,2,3], [4,5,6], [7,8,9]]))
indexs = tf.convert_to_tensor(np.array([1,1,1]))

@tf.function
def iterate_tensor(probs, indexs):
    return tf.linalg.diag_part(tf.gather(probs, indexs, axis=1))
iterate_tensor(probs, indexs)

Output:

How to iterate over a tf.tensor in eager mode

1 Answers