Используя вывод `tf.argmax` в качестве индекса в 3d

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?