description: TensorFlow variant of NumPy's take_along_axis.

tf.experimental.numpy.take_along_axis

View source on GitHub

TensorFlow variant of NumPy's take_along_axis.

tf.experimental.numpy.take_along_axis(
    arr, indices, axis
)

See the NumPy documentation for numpy.take_along_axis.