jax.numpy.diag

jax.numpy.diag(v, k=0)

返回指定的对角线或构造对角数组。

JAX 的 numpy.diag() 实现。

JAX 版本始终返回输入的副本，但如果在 JIT 编译中使用，编译器可能会避免复制。

参数:

• v (ArrayLike) – 输入数组。可以是 1 维数组来创建对角矩阵，也可以是 2 维数组来提取对角线。

• k (int) – 可选，默认值为0。对角线偏移量。正值将对角线置于主对角线上方，负值将对角线置于主对角线下方。

返回:

如果 v 是一个二维数组，则返回一个包含对角线元素的一维数组。如果 v 是一个一维数组，则返回一个二维数组，该数组的元素沿指定的对角线放置。

返回类型:

数组

另请参阅

• jax.numpy.diagflat()

• jax.numpy.diagonal()

示例

从一维数组创建对角矩阵

>>> jnp.diag(jnp.array([1, 2, 3]))
Array([[1, 0, 0],
       [0, 2, 0],
       [0, 0, 3]], dtype=int32)

指定对角线偏移量

>>> jnp.diag(jnp.array([1, 2, 3]), k=1)
Array([[0, 1, 0, 0],
       [0, 0, 2, 0],
       [0, 0, 0, 3],
       [0, 0, 0, 0]], dtype=int32)

从二维数组中提取对角线

>>> x = jnp.array([[1, 2, 3],
...                [4, 5, 6],
...                [7, 8, 9]])
>>> jnp.diag(x)
Array([1, 5, 9], dtype=int32)