Get all non-masked element of an array, and pad the rest

[Egiob]

Hello, I'm struggling on the following simple problem, I have an array of size $n$ with, say, $k$ non-masked elements and $n-k$ masked elements like this ($n=10$, $k=3$):

DeviceArray([ 0., nan, nan, nan, nan, nan, nan, 3., nan, 6.])

What I want is to get all non-nans elements and pad the rest, like this:

DeviceArray([ 0., 3., 6., nan, nan, nan, nan, nan, nan, nan])

I can assume my non-nans elements are sorted, and I could call jnp.sort, but this is quite an overkill for such simple problem (and it is slow).

Any idea how to efficiently perform this operation ? Any help is welcome, thanks in advance!

[jakevdp]

I think calling argsort on your filter condition is probably the simplest and fastest approach given the set of primitive operations JAX has to work with:

import jax.numpy as jnp
x = jnp.array([0., jnp.nan, jnp.nan, jnp.nan, 3., jnp.nan, 6.])
print(x[jnp.argsort(jnp.isnan(x))])
# [ 0.  3.  6. nan nan nan nan]

Since JAX uses a stable sort, it will keep your non-NaN values in their original order.

[YouJiacheng]

@Egiob Do you use jit in your test?

[Egiob]

I evaluate this way, with $x$ of size $1000$:

%timeit jnp.sort(x)
#66.4 µs ± 2.95 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each

%timeit x[jnp.nonzero(~jnp.isnan(x), size=10)]
#3.69 ms ± 83.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit jax.lax.sort_key_val(jnp.isnan(x), x)
#172 µs ± 4.41 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

[jakevdp]

Given the set of primitive operations available to JAX, I don't know of any fast way to do this that does not involve a sort operation. Another option would be to use jnp.where() with the size argument, but its JAX implementation relies on sorting under the hood.

PS - your benchmarks may give misleading results; for more info on how to microbenchmark JAX code, see https://jax.readthedocs.io/en/latest/faq.html#benchmarking-jax-code

[YouJiacheng]

@Egiob You should use jit, do warmup(compile) and block_until_ready.

f = jax.jit(lambda x: jnp.sort(x))
g = jax.jit(lambda x: x[jnp.nonzero(~jnp.isnan(x), size=10)])
h = jax.jit(lambda x: jax.lax.sort_key_val(jnp.isnan(x), x))
f(x), g(x), h(x) # compile
%timeit jax.block_until_ready(f(x))
%timeit jax.block_until_ready(g(x))
%timeit jax.block_until_ready(h(x))

[Egiob]

Thanks for your answers! My benchmark was misleading, indeed!

[YouJiacheng]

If k is a constant, you may use x[jnp.nonzero(~jnp.isnan(x), size=k)] then pad it to n elements.
Actually this is counting sort.

[Egiob]

It was the fastest solution for my problem, thanks a lot!