jax.numpy.nanmean#

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

傳回沿著指定軸的陣列元素平均值，並忽略 NaN 值。

NumPy v2.2 中 numpy.nanmean() 的 JAX 實作。

參數:

  • a (ArrayLike) – 輸入陣列。

  • axis (Axis) – int 或 int 序列，預設值=None。計算平均值的軸。如果為 None，則沿著展平的陣列計算平均值。

  • dtype (DTypeLike | None) – 輸出陣列的類型。預設值=None。

  • keepdims (bool) – bool，預設值=False。如果為 True，縮減的軸會保留在結果中，大小為 1。

  • where (ArrayLike | None) – 布林值類型的陣列，預設值=None。用於計算平均值的元素。陣列應與輸入廣播相容。

  • out (None) – JAX 未使用。

傳回值:

一個陣列，包含沿著指定軸的陣列元素平均值，並忽略 NaN 值。如果沿著指定軸的所有元素都是 NaN，則傳回 nan。

傳回類型:

Array

另請參閱

範例

預設情況下，jnp.nanmean 會計算沿著展平陣列的元素平均值。

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

如果 axis=1，則將沿著軸 1 計算平均值。

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

如果 keepdims=True，則輸出的 ndim 將與輸入的相同。

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

where 可用於僅包含特定元素來計算平均值。

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

如果 where 在所有元素處均為 False，則 jnp.nanmean 會沿著指定軸傳回 nan。

>>> where = jnp.array([[False],
...                    [False],
...                    [False]])
>>> jnp.nanmean(x, axis=0, keepdims=True, where=where)
Array([[nan, nan, nan, nan]], dtype=float32)