jax.numpy.tril_indices

jax.numpy.tril_indices#

jax.numpy.tril_indices(n, k=0, m=None)[源代码][源代码]#

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

JAX 实现的 numpy.tril_indices()。

参数:

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

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

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

返回:

包含两个数组的元组，每个数组分别包含沿每个轴的下三角形的索引。

返回类型:

tuple[Array, Array]

参见

示例

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

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

如果在输入中同时提供了 n 和 m，则返回 (n, m) 数组的下三角索引。

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

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

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

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

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