jax.numpy.tril_indices_from

jax.numpy.tril_indices_from#

jax.numpy.tril_indices_from(arr, k=0)[源代码][源代码]#

返回给定数组的下三角索引。

JAX 实现的 numpy.tril_indices_from()。

参数:

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

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

返回:

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

返回类型:

tuple[Array, Array]

参见

示例

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

由 jnp.tril_indices_from 索引的元素对应于 jnp.tril 输出的元素。

>>> ind = jnp.tril_indices_from(arr)
>>> arr[ind]
Array([1, 4, 5, 7, 8, 9], dtype=int32)
>>> jnp.tril(arr)
Array([[1, 0, 0],
       [4, 5, 0],
       [7, 8, 9]], dtype=int32)

当 k > 0 时：

>>> jnp.tril_indices_from(arr, 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 < 0 时：

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