tf.keras.ops.take

Take elements from a tensor along an axis.

xSource tensor.
indicesThe indices of the values to extract.
axisThe axis over which to select values. By default, the flattened input tensor is used.

The corresponding tensor of values.