jax.numpy.tri

jax.numpy.tri(N, M=None, k=0, dtype=None)

返回一个在对角线及其下方为 1，其他地方为 0 的数组。

numpy.tri() 的 JAX 实现

参数:

• N (int) – int. 返回数组的行维度。

• M (int | None | None) – 可选，int. 返回数组的列维度。 如果未指定，则 M = N。

• k (int) – 可选，整数，默认值=0。指定数组中填充 1 的子对角线及其下方的位置。k=0 指的是主对角线，k<0 指的是主对角线之下的子对角线，而 k>0 指的是主对角线之上的子对角线。

• dtype (DTypeLike | None | None) – 可选，返回数组的数据类型。默认类型为浮点型。

返回:

一个形状为 (N, M) 的数组，其中由 k 指定的子对角线之下的元素设置为 1，其他位置设置为 0。

返回类型:

数组

另请参阅

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

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

示例

>>> jnp.tri(3)
Array([[1., 0., 0.],
       [1., 1., 0.],
       [1., 1., 1.]], dtype=float32)

当 M 不等于 N 时

>>> jnp.tri(3, 4)
Array([[1., 0., 0., 0.],
       [1., 1., 0., 0.],
       [1., 1., 1., 0.]], dtype=float32)

当 k>0 时

>>> jnp.tri(3, k=1)
Array([[1., 1., 0.],
       [1., 1., 1.],
       [1., 1., 1.]], dtype=float32)

当 k<0 时

>>> jnp.tri(3, 4, k=-1)
Array([[0., 0., 0., 0.],
       [1., 0., 0., 0.],
       [1., 1., 0., 0.]], dtype=float32)