jax.numpy.diag_indices_from

jax.numpy.diag_indices_from(arr)

返回用于访问给定数组主对角线的索引。

JAX 对 numpy.diag_indices_from() 的实现。

参数:

arr (ArrayLike) – 输入数组。必须至少是 2 维的，并且在所有维度上具有相等的长度。

返回值:

一个数组元组，包含用于访问输入数组主对角线的索引。

返回类型:

tuple[Array, …]

另请参阅

• jax.numpy.diag_indices()

• jax.numpy.diagonal()

示例

>>> arr = jnp.array([[1, 2, 3],
...                  [4, 5, 6],
...                  [7, 8, 9]])
>>> jnp.diag_indices_from(arr)
(Array([0, 1, 2], dtype=int32), Array([0, 1, 2], dtype=int32))
>>> arr = jnp.array([[[1, 2], [3, 4]],
...                  [[5, 6], [7, 8]]])
>>> jnp.diag_indices_from(arr)
(Array([0, 1], dtype=int32),
Array([0, 1], dtype=int32),
Array([0, 1], dtype=int32))