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