cr.nimble.dsp.largest_indices_rw

cr.nimble.dsp.largest_indices_rw(X, K)[source]

Returns the indices of K largest entries by magnitude in each row of X

Parameters

  • X (jax.numpy.ndarray) – An (S,N) data matrix with data points in rows

  • K (int) – The number of largest entries to be identified in each row of X

Returns

An (S,K) index matrix indices of K largest elements in each row of X

Return type

(jax.numpy.ndarray)