tf.keras.ops.take_along_axis

Select values from x at the 1-D indices along the given axis.

View aliases

Main aliases

tf.keras.ops.numpy.take_along_axis

tf.keras.ops.take_along_axis(
    x, indices, axis=None
)

Args
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.

Returns
The corresponding tensor of values.