cr.nimble.dsp.take_along_cols

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

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

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

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

Returns

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

Return type

(jax.numpy.ndarray)