JaxBackend.topk_indices

JaxBackend.topk_indices(arr, k)

Determinces the indices of largest elements.

Parameters:

arrBackendArray

Input data.

kint

Number of maxima to determine.

Returns:

BackendArray

Indices of k largest elements in arr.