mlx.core.topk

topk(a: array, /, k: int, axis: None | int = -1, *, stream: None | Stream | Device = None) → array

Returns the k largest elements from the input along a given axis.

The elements will not necessarily be in sorted order.

参数:

• a (array) -- Input array.

• k (int) -- k top elements to be returned

• axis (int or None, optional) -- Optional axis to select over. If None, this selects the top k elements over the flattened array. If unspecified, it defaults to -1.

返回:

The top k elements from the input.

返回类型:

array