tf.keras.ops.argpartition

Performs an indirect partition along the given axis.

View aliases

Main aliases

tf.keras.ops.numpy.argpartition

tf.keras.ops.argpartition(
    x, kth, axis=-1
)

It returns an array of indices of the same shape as x that index data along the given axis in partitioned order.

Args
x	Array to sort.
kth	Element index to partition by. The k-th element will be in its final sorted position and all smaller elements will be moved before it and all larger elements behind it. The order of all elements in the partitions is undefined. If provided with a sequence of k-th it will partition all of them into their sorted position at once.
axis	Axis along which to sort. The default is -1 (the last axis). If None, the flattened array is used.

Returns
Array of indices that partition x along the specified axis.