jax.numpy.diag

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

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

numpy.diag() 的 JAX 实现。

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

参数:

• v (ArrayLike) – 输入数组。可以是创建对角矩阵的一维数组或提取对角线的二维数组。

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