jax.numpy.tril

jax.numpy.tril(m, k=0)

返回数组的下三角部分。

JAX 对 numpy.tril() 的实现

参数:

• m (类数组) – 输入数组。 必须满足 m.ndim >= 2。

• k (int) – k：可选，整数，默认为0。 指定数组元素设置为零的子对角线。 k=0 指的是主对角线，k<0 指的是主对角线以下的子对角线，k>0 指的是主对角线以上的子对角线。

返回:

一个与输入具有相同形状的数组，其中包含给定数组的下三角部分，并且由 k 指定的子对角线以下的元素设置为零。

返回类型:

数组

另请参阅

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

• jax.numpy.tri(): 返回一个在对角线及其下方为 1，其他地方为 0 的数组。

示例

>>> x = jnp.array([[1, 2, 3, 4],
...                [5, 6, 7, 8],
...                [9, 10, 11, 12]])
>>> jnp.tril(x)
Array([[ 1,  0,  0,  0],
       [ 5,  6,  0,  0],
       [ 9, 10, 11,  0]], dtype=int32)
>>> jnp.tril(x, k=1)
Array([[ 1,  2,  0,  0],
       [ 5,  6,  7,  0],
       [ 9, 10, 11, 12]], dtype=int32)
>>> jnp.tril(x, k=-1)
Array([[ 0,  0,  0,  0],
       [ 5,  0,  0,  0],
       [ 9, 10,  0,  0]], dtype=int32)

当 m.ndim > 2 时， jnp.tril 对尾部轴进行批量操作。

>>> x1 = jnp.array([[[1, 2],
...                  [3, 4]],
...                 [[5, 6],
...                  [7, 8]]])
>>> jnp.tril(x1)
Array([[[1, 0],
        [3, 4]],

       [[5, 0],
        [7, 8]]], dtype=int32)