jax.numpy.median

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

返回沿给定轴的数组元素的中位数。

JAX 实现 numpy.median().

参数:

• a (ArrayLike) – 输入数组。

• axis (int | tuple[int, ...] | None) – 可选，int 或 int 序列，默认值为 None。要计算中位数的轴。如果为 None，则计算扁平化数组的中位数。

• keepdims (bool) – bool，默认值为 False。如果为 True，则保留结果中大小为 1 的缩减轴。

• out (None) – JAX 不使用。

• overwrite_input (bool) – JAX 不使用。

返回:

沿给定轴的中位数数组。

返回类型:

数组

另请参见

• jax.numpy.mean(): 计算给定轴上数组元素的平均值。

• jax.numpy.max(): 计算给定轴上数组元素的最大值。

• jax.numpy.min(): 计算给定轴上数组元素的最小值。

示例

默认情况下，计算扁平化数组的中位数。

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

如果 axis=1，则沿轴 1 计算中位数。

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

如果 keepdims=True，则输出的 ndim 等于输入的 ndim。

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