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是二维数组，返回包含对角线元素的一维数组。
示例
- 创建对角矩阵
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]]
- 提取对角线元素
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]
- 提取偏移对角线元素
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 是一个非常有用的函数，可以根据需要创建对角矩阵或提取矩阵的对角线元素。在线性代数和矩阵计算中，这个功能常用于各种数值计算和矩阵操作中。