jax.numpy.take

目录

jax.numpy.take#

jax.numpy.take(a, indices, axis=None, out=None, mode=None, unique_indices=False, indices_are_sorted=False, fill_value=None)[源代码][源代码]#

从数组中提取元素。

JAX 实现的 numpy.take()，基于 jax.lax.gather() 实现。JAX 的行为在处理越界索引时与 NumPy 不同；请参见下面的 mode 参数。

参数:

  • a (ArrayLike) – 从中获取值的数组。

  • indices (ArrayLike) – 从数组中获取值的整数索引的N维数组。

  • axis (int | None) – 要沿其取值的轴。如果未指定，将在应用索引之前将数组展平。

  • mode (str | None) – 越界索引模式，可以是 "fill" 或 "clip"。默认的 mode="fill" 对越界索引返回无效值（例如 NaN）；fill_value 参数控制此值。有关 mode 选项的更多讨论，请参见 jax.numpy.ndarray.at。

  • fill_value (StaticScalar | None) – 当模式为 ‘fill’ 时，返回超出边界的切片所使用的填充值。否则忽略。对于不精确类型，默认值为 NaN；对于有符号类型，默认值为最大负值；对于无符号类型，默认值为最大正值；对于布尔类型，默认值为 True。

  • unique_indices (bool) – 如果为True，实现将假设索引是唯一的，这可以在某些后端上实现更高效的执行。如果设置为True且索引不唯一，输出结果将不确定。

  • indices_are_sorted (bool) – 如果为 True，实现将假设索引按升序排序，这可以在某些后端上实现更高效的执行。如果设置为 True 且索引未排序，则输出未定义。

  • out (None)

返回:

从 a 中提取的值数组。

返回类型:

Array

参见

示例

>>> x = jnp.array([[1., 2., 3.],
...                [4., 5., 6.]])
>>> indices = jnp.array([2, 0])

不传递轴参数会导致对展平数组进行索引：

>>> jnp.take(x, indices)
Array([3., 1.], dtype=float32)
>>> x.ravel()[indices]  # equivalent indexing syntax
Array([3., 1.], dtype=float32)

传递一个轴会导致将索引应用于沿该轴的每个子数组：

>>> jnp.take(x, indices, axis=1)
Array([[3., 1.],
       [6., 4.]], dtype=float32)
>>> x[:, indices]  # equivalent indexing syntax
Array([[3., 1.],
       [6., 4.]], dtype=float32)

越界索引填充为无效值。对于浮点输入，这是 NaN：

>>> jnp.take(x, indices, axis=0)
Array([[nan, nan, nan],
       [ 1.,  2.,  3.]], dtype=float32)
>>> x.at[indices].get(mode='fill', fill_value=jnp.nan)  # equivalent indexing syntax
Array([[nan, nan, nan],
       [ 1.,  2.,  3.]], dtype=float32)

可以使用 mode 参数调整这种默认的越界行为，例如，我们可以改为裁剪到最后一个有效值：

>>> jnp.take(x, indices, axis=0, mode='clip')
Array([[4., 5., 6.],
       [1., 2., 3.]], dtype=float32)
>>> x.at[indices].get(mode='clip')  # equivalent indexing syntax
Array([[4., 5., 6.],
       [1., 2., 3.]], dtype=float32)
目录