jax.numpy.nanmedian

jax.numpy.nanmedian(a, axis=None, out=None, overwrite_input=False, keepdims=False)

Return the median of array elements along a given axis, ignoring NaNs.

JAX implementation of numpy.nanmedian().

Parameters:

• a (ArrayLike) – input array.

• axis (int | tuple[int, ...] | None) – optional, int or sequence of ints, default=None. Axis along which the median to be computed. If None, median is computed for the flattened array.

• keepdims (bool) – bool, default=False. If true, reduced axes are left in the result with size 1.

• out (None) – Unused by JAX.

• overwrite_input (bool) – Unused by JAX.

Returns:

An array containing the median along the given axis, ignoring NaNs. If all elements along the given axis are NaNs, returns nan.

Return type:

Array

See also

• jax.numpy.nanmean(): Compute the mean of array elements over a given axis, ignoring NaNs.

• jax.numpy.nanmax(): Compute the maximum of array elements over given axis, ignoring NaNs.

• jax.numpy.nanmin(): Compute the minimum of array elements over given axis, ignoring NaNs.

Examples

By default, the median is computed for the flattened array.

>>> nan = jnp.nan
>>> x = jnp.array([[2, nan, 7, nan],
...                [nan, 5, 9, 2],
...                [6, 1, nan, 3]])
>>> jnp.nanmedian(x)
Array(4., dtype=float32)

If axis=1, the median is computed along axis 1.

>>> jnp.nanmedian(x, axis=1)
Array([4.5, 5. , 3. ], dtype=float32)

If keepdims=True, ndim of the output is equal to that of the input.

>>> jnp.nanmedian(x, axis=1, keepdims=True)
Array([[4.5],
       [5. ],
       [3. ]], dtype=float32)