jax.numpy.triu_indices_from

jax.numpy.triu_indices_from(arr, k=0)

返回给定数组上三角部分的索引。

JAX 对 numpy.triu_indices_from() 的实现。

参数:

• arr (ArrayLike | SupportsShape) – 输入数组。必须满足 arr.ndim == 2。

• k (int) – 可选，整数，默认值为 0。指定了上三角部分的索引所在的对角线及其上方的对角线。 k=0 指的是主对角线，k<0 指的是主对角线以下的次对角线，而 k>0 指的是主对角线上方的次对角线。

返回:

包含上三角部分索引的两个数组组成的元组，每个轴一个。

返回类型:

tuple[Array, Array]

另请参阅

• jax.numpy.tril_indices_from()：返回给定数组的下三角部分的索引。

• jax.numpy.triu_indices()：返回大小为 (n, m) 的数组的上三角部分的索引。

• jax.numpy.triu()：返回数组的上三角部分。

示例

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

通过 jnp.triu_indices_from 索引的元素对应于 jnp.triu 的输出中的元素。

>>> ind = jnp.triu_indices_from(arr)
>>> arr[ind]
Array([1, 2, 3, 5, 6, 9], dtype=int32)
>>> jnp.triu(arr)
Array([[1, 2, 3],
       [0, 5, 6],
       [0, 0, 9]], dtype=int32)

当 k > 0 时

>>> jnp.triu_indices_from(arr, k=1)
(Array([0, 0, 1], dtype=int32), Array([1, 2, 2], dtype=int32))

当 k < 0 时

>>> jnp.triu_indices_from(arr, k=-1)
(Array([0, 0, 0, 1, 1, 1, 2, 2], dtype=int32), Array([0, 1, 2, 0, 1, 2, 1, 2], dtype=int32))