jax.numpy.triu_indices

jax.numpy.triu_indices(n, k=0, m=None)

返回大小为 (n, m) 的数组的上三角部分的索引。

JAX 对 numpy.triu_indices() 的实现。

参数:

• n (DimSize) – int。要返回索引的数组的行数。

• k (DimSize) – 可选，int，默认为 0。指定要返回上三角索引的对角线以及其上方的对角线。 k=0 指的是主对角线， k<0 指的是主对角线下方子对角线， k>0 指的是主对角线上方子对角线。

• m (DimSize | None) – 可选，int。要返回索引的数组的列数。如果未指定，则 m = n。

返回:

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

返回类型:

tuple[Array, Array]

另请参阅

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

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

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

示例

如果输入中只提供了 n，则返回大小为 (n, n) 的数组的上三角部分的索引。

>>> jnp.triu_indices(3)
(Array([0, 0, 0, 1, 1, 2], dtype=int32), Array([0, 1, 2, 1, 2, 2], dtype=int32))

如果输入中同时提供了 n 和 m，则返回大小为 (n, m) 的数组的上三角部分的索引。

>>> jnp.triu_indices(3, m=2)
(Array([0, 0, 1], dtype=int32), Array([0, 1, 1], dtype=int32))

如果 k = 1，则返回主对角线上方第一个子对角线及其上方的索引。

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

如果 k = -1，则返回主对角线下第一个子对角线及其上方的索引。

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