Что является эквивалентом этого кода Numpy в TensorFlow?

Question

Я манипулирую массивами Numpy, и код выглядит следующим образом:

z[np.arange(n), y]

Где z - двумерный массив, у - одномерный массив. Кроме того, z.shape [0] == y.shape [0] == n.

Как я могу сделать эквивалентные вещи с тензорами TensorFlow?

Answer 1

Вы можете использовать tf.gather_nd, чтобы получить нужную индексацию.

import numpy as np
import tensorflow as tf

# Numpy implementation
n = 3
z = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
y = np.array([0, 1, 1])

assert z.shape[0] == y.shape[0] == n

np_out = z[np.arange(n), y]

# TF implementation
tf.reset_default_graph()

range_t = tf.range(n) # Equiv to np.arange
x_y = tf.stack([range_t, y], axis=1) # Get (x,y) as a tuple
pick_by_index_from_z = tf.gather_nd(z, x_y) # Pick the right values from z

with tf.Session() as sess:
  tf_out = sess.run(pick_by_index_from_z)

# The np and tf values should be the same
assert (np_out == tf_out).all()

print('z:')
print(z)
print('\nnp_out:')
print(np_out)
print('\ntf_out:')
print(tf_out)

Это дает вывод:

z:
[[1 2 3]
 [4 5 6]
 [7 8 9]]

np_out:
[1 5 8]

tf_out:
[1 5 8]