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