cr.nimble.dsp.take_along_rows

cr.nimble.dsp.take_along_rows(X, indices)[source]

Picks K entries from each row of X specified by indices matrix

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

  • indices (jax.numpy.ndarray) – An (S,K) index matrix identifying the values to be picked up from X

Returns

An (S,K) data matrix subset of X containing K elements from each row of X

Return type

(jax.numpy.ndarray)