jax.numpy.nanargmin

jax.numpy.nanargmin(a, axis=None, out=None, keepdims=None)

返回数组中最小值的索引，忽略 NaN。

JAX 实现 numpy.nanargmin()。

参数:

• a (ArrayLike) – 输入数组

• axis (int | None) – 可选的整数，指定查找最大值所在的轴。如果未指定 axis，则 a 将被展平。

• out (None) – JAX 未使用

• keepdims (bool | None) – 如果为 True，则返回一个与 a 具有相同维度的数组。

返回:

包含沿指定轴的最小值的索引的数组。

返回类型:

Array

注意

如果某个轴包含所有 NaN 值，则返回的索引将为 -1。这与 numpy.nanargmin() 抛出错误的行为不同。

另请参阅

• jax.numpy.argmin()：返回最小值的索引。

• jax.numpy.nanargmax()：在忽略 NaN 值的情况下计算 argmax。

示例

>>> x = jnp.array([jnp.nan, 3, 5, 4, 2])
>>> jnp.nanargmin(x)
Array(4, dtype=int32)

>>> x = jnp.array([[1, 3, jnp.nan],
...                [5, 4, jnp.nan]])
>>> jnp.nanargmin(x, axis=1)
Array([0, 1], dtype=int32)

>>> jnp.nanargmin(x, axis=1, keepdims=True)
Array([[0],
       [1]], dtype=int32)