jax.numpy.ravel_multi_index

jax.numpy.ravel_multi_index#

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

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

JAX 实现 numpy.ravel_multi_index()

参数：

  • 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 样式列优先顺序。

返回值:

扁平化索引的数组

返回值类型:

数组

另请参见

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