在3D中使用`tf.argmax`的输出作为索引

Question

我想使用'tf.argmax'的输出作为3d中的索引来获取值。 这是下面的示例代码。

import tensorflow as tf
inp = tf.constant([[[1, 2, 3, 4, 5], [1, 2, 6, 4, 5], [1, 2, 6, 4, 5]], [[1, 3, 4, 4, 1], [1, 2, 1, 4, 5], [1, 2, 6, 4, 5]]])
sess = tf.Session()

res = tf.argmax(inp, axis=2)

x = tf.constant([[[4, 2, 3, 4, 5], [1, 11, 6, 4, 5], [1, 5, 6, 4, 5]], [[1, 3, 1, 4, 1], [1, 2, 1, 4, 1], [1, 2, 2, 4, 2]]])

print(sess.run(x [:, :, res]))

如何使用索引从3d中的其他张量获取值？