jax.numpy.diag

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

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

JAX 实现的 numpy.diag()。

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

参数:

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

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

返回:

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

返回类型:

Array

参见

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