jax.numpy.median#

jax.numpy.median(a, axis=None, out=None, overwrite_input=False, keepdims=False)[source]#

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

numpy.median() 的 JAX 实现。

参数:
  • a (ArrayLike) – 输入数组。

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

  • keepdims (bool) – 布尔值,默认值=False。如果为 true,则缩减的轴将保留在结果中,大小为 1。

  • out (None) – JAX 未使用。

  • overwrite_input (bool) – JAX 未使用。

返回值:

沿给定轴的中位数数组。

返回类型:

Array

另请参阅

示例

默认情况下,为扁平化数组计算中位数。

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