Return the mean of the array elements along a given axis, ignoring NaNs.
JAX implementation of numpy.nanmean()
.
a (ArrayLike) – Input array.
axis (Axis) – int or sequence of ints, default=None. Axis along which the mean is computed. If None, the mean is computed along the flattened array.
dtype (DTypeLike | None) – The type of the output array. Default=None.
keepdims (bool) – bool, default=False. If True, reduced axes are left in the result with size 1.
where (ArrayLike | None) – array of boolean dtype, default=None. The elements to be used in computing mean. Array should be broadcast compatible to the input.
out (None) – Unused by JAX.
An array containing the mean of array elements along the given axis, ignoring NaNs. If all elements along the given axis are NaNs, returns nan
.
See also
jax.numpy.nanmin()
: Compute the minimum of array elements along a given axis, ignoring NaNs.
jax.numpy.nanmax()
: Compute the maximum of array elements along a given axis, ignoring NaNs.
jax.numpy.nansum()
: Compute the sum of array elements along a given axis, ignoring NaNs.
jax.numpy.nanprod()
: Compute the product of array elements along a given axis, ignoring NaNs.
Examples
By default, jnp.nanmean
computes the mean of elements along the flattened array.
>>> nan = jnp.nan >>> x = jnp.array([[2, nan, 4, 3], ... [nan, -2, nan, 9], ... [4, -7, 6, nan]]) >>> jnp.nanmean(x) Array(2.375, dtype=float32)
If axis=1
, mean will be computed along axis 1.
>>> jnp.nanmean(x, axis=1) Array([3. , 3.5, 1. ], dtype=float32)
If keepdims=True
, ndim
of the output will be same of that of the input.
>>> jnp.nanmean(x, axis=1, keepdims=True) Array([[3. ], [3.5], [1. ]], dtype=float32)
where
can be used to include only specific elements in computing the mean.
>>> where = jnp.array([[1, 0, 1, 0], ... [0, 0, 1, 1], ... [1, 1, 0, 1]], dtype=bool) >>> jnp.nanmean(x, axis=1, keepdims=True, where=where) Array([[ 3. ], [ 9. ], [-1.5]], dtype=float32)
If where
is False
at all elements, jnp.nanmean
returns nan
along the given axis.
>>> where = jnp.array([[False], ... [False], ... [False]]) >>> jnp.nanmean(x, axis=0, keepdims=True, where=where) Array([[nan, nan, nan, nan]], dtype=float32)
RetroSearch is an open source project built by @garambo | Open a GitHub Issue
Search and Browse the WWW like it's 1997 | Search results from DuckDuckGo
HTML:
3.2
| Encoding:
UTF-8
| Version:
0.7.3