tensorflow - how to get value from a index value

Question

The tensor is A_tensor,has shape [3,3,3] and I want get the last axis's value by index.

How do I do that in Tensorflow?

A_tensor =tf.constant([[1,2,3],[2,3,4],[3,4,5]])

How to get tensor ([[1,2],[2,3],[3,4]])?

question from:https://stackoverflow.com/questions/65622728/how-to-get-value-from-a-index-value

Answer 1

You can use tf.gather():

A_tensor =tf.constant([[1,2,3],[2,3,4],[3,4,5]])
tf.gather(A_tensor,[[0,1]],axis=-1)