jax.numpy.tri

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

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

JAX 实现 numpy.tri()

参数:

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

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

• k (int) – 可选，int，默认值=0。指定填充 1 的子对角线及其下方的对角线。 k=0 指主对角线，k<0 指主对角线下方的子对角线，k>0 指主对角线上方的子对角线。

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

返回:

一个形状为 (N, M) 的数组，其中在由 k 指定的子对角线及其以下的元素设置为一，其余元素为零。

返回类型:

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)