jax.numpy.nanmean#

jax.numpy.nanmean(a,axis=None,dtype=None,out=None,keepdims=False,where=None)[source]#

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

JAX implementation ofnumpy.nanmean().

Parameters:

• a (ArrayLike) – Input array.

• axis (Axis) – int or sequence of ints, default=None. Axis along which the mean iscomputed. 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 resultwith size 1.

• where (ArrayLike |None) – array of boolean dtype, default=None. The elements to be used incomputing mean. Array should be broadcast compatible to the input.

• out (None) – Unused by JAX.

Returns:

An array containing the mean of array elements along the given axis, ignoringNaNs. If all elements along the given axis are NaNs, returnsnan.

Return type:

Array

See also

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

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

• jax.numpy.nansum(): Compute the sum of array elements along a givenaxis, ignoring NaNs.

• jax.numpy.nanprod(): Compute the product of array elements along agiven axis, ignoring NaNs.

Examples

By default,jnp.nanmean computes the mean of elements along the flattenedarray.

>>>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)

Ifaxis=1, mean will be computed along axis 1.

>>>jnp.nanmean(x,axis=1)Array([3. , 3.5, 1. ], dtype=float32)

Ifkeepdims=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)

Ifwhere isFalse at all elements,jnp.nanmean returnsnanalong 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)