jax.numpy.triu_indices

jax.numpy.triu_indices#

jax.numpy.triu_indices(n, k=0, m=None)[source]#

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

JAX 实现的 numpy.triu_indices()。

参数：

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

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

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

返回值：

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

返回类型：

tuple[Array, Array]

另请参阅

示例

如果输入中只提供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))
内容