jax.scipy.linalg.sqrtm

目录

jax.scipy.linalg.sqrtm#

jax.scipy.linalg.sqrtm(A, blocksize=1)[源代码][源代码]#

计算矩阵平方根

JAX 实现的 scipy.linalg.sqrtm()。

参数:

  • A (ArrayLike) – 形状为 (N, N) 的数组

  • blocksize (int) – 在JAX中不支持；JAX始终使用 blocksize=1。

返回:

形状为 (N, N) 的数组，包含矩阵 A 的平方根

返回类型:

Array

参见

jax.scipy.linalg.expm()

示例

>>> a = jnp.array([[1., 2., 3.],
...                [2., 4., 2.],
...                [3., 2., 1.]])
>>> sqrt_a = jax.scipy.linalg.sqrtm(a)
>>> with jnp.printoptions(precision=2, suppress=True):
...   print(sqrt_a)
[[0.92+0.71j 0.54+0.j   0.92-0.71j]
 [0.54+0.j   1.85+0.j   0.54-0.j  ]
 [0.92-0.71j 0.54-0.j   0.92+0.71j]]

根据定义，矩阵平方根与其自身的矩阵乘法应等于输入：

>>> jnp.allclose(a, sqrt_a @ sqrt_a)
Array(True, dtype=bool)

备注

此函数实现了 [1] 中描述的复杂Schur方法。由于JAX中尚无Sylvester方程求解器，因此它不使用递归阻塞来加速计算。

引用

目录