jax.numpy.diag_indices_from

jax.numpy.diag_indices_from(arr)

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

JAX 实现的 numpy.diag_indices_from()。

参数:

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

返回:

包含访问输入数组主对角线索引的数组元组。

返回类型:

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