jax.numpy.tri

内容

jax.numpy.tri#

jax.numpy.tri(N, M=None, k=0, dtype=None)[source]#

返回一个数组，在对角线及其下方为 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) – 可选，返回数组的数据类型。默认类型为 float。

返回值：

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

返回类型：

数组

另请参阅

示例

>>> 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)
内容