jax.numpy.diag#
- jax.numpy.diag(v, k=0)[源码]#
返回指定的对角线或构造对角线数组。
JAX 实现的
numpy.diag()
。JAX 版本始终返回输入的副本,尽管如果在 JIT 编译中使用,编译器可能会避免复制。
- 参数:
v (ArrayLike) – 输入数组。可以是一维数组以创建对角矩阵,也可以是二维数组以提取对角线。
k (int) – 可选,默认=0。对角线偏移。正值将对角线置于主对角线上方,负值将其置于主对角线下方。
- 返回:
如果 v 是二维数组,则返回包含对角线元素的一维数组。如果 v 是一维数组,则返回一个二维数组,其中输入元素沿指定的对角线放置。
- 返回类型:
示例
从一维数组创建对角矩阵
>>> 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)