Clean masked median?

[jeffgortmaker]

I'm having some trouble coming up with a clean way of computing a median over a (dynamic) mask in a jitted function. Means are simple because of the where argument: jnp.mean(x, where=mask). But NumPy and hence JAX don't support where for jnp.median.

The mask has a dynamic number of True elements, so I don't think something like jnp.median(x.at[mask].get()) will work in my setting. The best I've come up with so far is jnp.nanmedian(jnp.where(mask, x, jnp.nan)), but this has the undesirable side-effect of always triggering the jax_debug_nans config setting, which I like to keep on when debugging other parts of my code.

I'm wondering if there's a cleaner way of doing this with JAX that I'm not thinking of?

[jakevdp]

Using nanmedian seems the cleanest, but if you want to avoid introducing NaNs you could do something similar to the implementation of nanquantile. It might look more-or-less like this:

import jax.numpy as jnp
import numpy as np
import jax

@jax.jit
def median_where(x, where):
  assert x.ndim == 1
  assert x.shape == where.shape
  assert where.dtype == bool
  N = jnp.sum(where)
  x = jnp.where(where, x, x.max())
  return jnp.quantile(x, 0.5 * N / len(x), method='midpoint')

np.random.seed(42)
x = jnp.array(np.random.randn(100))
where = (x > 0)
print(jnp.median(x[where]))
# 0.52791375
print(median_where(x, where))
# 0.52791375

[jeffgortmaker]

Thanks! That's great.