jax.numpy.tri

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

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

JAX 实现的 numpy.tri()

参数:

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

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

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

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

返回:

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

返回类型:

Array

参见

• 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)