jax.numpy.ravel_multi_index#

jax.numpy.ravel_multi_index(multi_index, dims, mode='raise', order='C')[源代码]#

将多维索引转换为扁平索引。

numpy.ravel_multi_index()的JAX实现

参数:

  • multi_index (Sequence[ArrayLike]) – 包含每个维度索引的整数数组序列。

  • dims (Sequence[int]) – 整数大小的序列; 必须满足 len(dims) == len(multi_index)

  • mode (str) –

    如何处理越界索引。选项包括

    • "raise" (默认): 引发 ValueError。此模式与 jit() 或其他 JAX 转换不兼容。

    • "clip": 将越界索引裁剪到有效范围。

    • "wrap": 将越界索引包裹到有效范围。

  • order (str) – "C" (默认) 或 "F"，指定是假定 C 风格的行优先顺序还是 Fortran 风格的列优先顺序。

返回:

扁平化索引的数组

返回类型:

Array

另请参阅

jax.numpy.unravel_index(): 此函数的逆函数。

示例

定义一个二维数组和偶数值的索引序列

>>> x = jnp.array([[2., 3., 4.],
...                [5., 6., 7.]])
>>> indices = jnp.where(x % 2 == 0)
>>> indices
(Array([0, 0, 1], dtype=int32), Array([0, 2, 1], dtype=int32))
>>> x[indices]
Array([2., 4., 6.], dtype=float32)

计算扁平化索引

>>> indices_flat = jnp.ravel_multi_index(indices, x.shape)
>>> indices_flat
Array([0, 2, 4], dtype=int32)

这些扁平化索引可用于从扁平化的 x 数组中提取相同的值

>>> x_flat = x.ravel()
>>> x_flat
Array([2., 3., 4., 5., 6., 7.], dtype=float32)
>>> x_flat[indices_flat]
Array([2., 4., 6.], dtype=float32)

可以使用 unravel_index() 恢复原始索引

>>> jnp.unravel_index(indices_flat, x.shape)
(Array([0, 0, 1], dtype=int32), Array([0, 2, 1], dtype=int32))