interpax_fft.take_mask

interpax_fft.take_mask(a, mask, /, *, size=-1, fill_value=None)Source

JIT compilable method to return a[mask][:size] padded by fill_value.

Parameters:

  • a (jnp.ndarray) – The source array.

  • mask (jnp.ndarray) – Boolean mask to index into a. Should have same shape as a.

  • size (int) – Elements of a at the first size True indices of mask will be returned. Size is clipped to be <= mask.size. A negative value will be interpreted as mask.size.

  • fill_value (Any) – When there are fewer nonzero elements in mask than size, the remaining elements will be filled with fill_value. Defaults to NaN for inexact types, the largest negative value for signed types, the largest positive value for unsigned types, and True for booleans.

Returns:

result (jnp.ndarray) – Shape (size, ).