jnp.diag 是 JAX 库中用于创建对角矩阵或提取对角线元素的函数。具体功能取决于输入的形状：

• 当输入是一维数组时，jnp.diag 创建一个以该数组为对角线元素的对角矩阵。
• 当输入是二维数组时，jnp.diag 提取并返回对角线元素。

函数签名

jnp.diag(v, k=0)

参数说明

• v (array_like): 输入数组。一维数组时，jnp.diag 创建一个对角矩阵；二维数组时，jnp.diag 提取并返回对角线元素。
• k (int, optional): 对角线的索引。k=0 表示主对角线，k>0 表示上方的偏移对角线，k<0 表示下方的偏移对角线。默认值为 0。

返回值

• 如果 v 是一维数组，返回一个二维对角矩阵。
• 如果 v 是二维数组，返回包含对角线元素的一维数组。

示例

1. 创建对角矩阵

import jax.numpy as jnp# 一维数组
v = jnp.array([1, 2, 3])# 创建对角矩阵
diag_matrix = jnp.diag(v)print("对角矩阵:\n", diag_matrix)

输出：

对角矩阵:[[1 0 0][0 2 0][0 0 3]]

2. 提取对角线元素

import jax.numpy as jnp# 二维数组
A = jnp.array([[1, 2, 3],[4, 5, 6],[7, 8, 9]])# 提取主对角线元素
diag_elements = jnp.diag(A)print("主对角线元素:", diag_elements)

输出：

主对角线元素: [1 5 9]

3. 提取偏移对角线元素

import jax.numpy as jnp# 二维数组
A = jnp.array([[1, 2, 3],[4, 5, 6],[7, 8, 9]])# 提取上方偏移对角线元素 (k=1)
diag_elements_upper = jnp.diag(A, k=1)print("上方偏移对角线元素:", diag_elements_upper)# 提取下方偏移对角线元素 (k=-1)
diag_elements_lower = jnp.diag(A, k=-1)print("下方偏移对角线元素:", diag_elements_lower)

输出：

上方偏移对角线元素: [2 6]
下方偏移对角线元素: [4 8]

总结

jnp.diag 是一个非常有用的函数，可以根据需要创建对角矩阵或提取矩阵的对角线元素。在线性代数和矩阵计算中，这个功能常用于各种数值计算和矩阵操作中。