jax.numpy.diag_indices_from#
- jax.numpy.diag_indices_from(arr)[源代码]#
返回用于访问给定数组主对角线的索引。
numpy.diag_indices_from()
的 JAX 实现。- 参数:
arr (ArrayLike) – 输入数组。必须至少为 2 维,并且所有维度上的长度相等。
- 返回:
一个数组元组,其中包含访问输入数组主对角线的索引。
- 返回类型:
示例
>>> 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))