tf.keras.ops.take

Take elements from a tensor along an axis.

x Source tensor.
indices The indices of the values to extract.
axis The axis over which to select values. By default, the flattened input tensor is used.

The corresponding tensor of values.