jax.scipy.linalg.funm#

jax.scipy.linalg.funm(A, func, disp=True)[源代码]#

计算矩阵值函数

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

参数:

  • A (ArrayLike) – 要计算函数的形状为 (N, N) 的数组。

  • func (Callable[[Array], Array]) – 可调用对象，接受一个标量参数并返回一个标量结果。表示要在 A 的特征值上计算的函数。

  • disp (bool) – 如果为 true（默认值），则不返回错误信息。与 scipy 版本不同，JAX 不会在运行时尝试显示信息。

  • compute_expm – (N, N) 类数组或 None（可选）。如果提供，则为 A 的矩阵指数。当 func 为指数函数时，这用于提高效率。如果未提供，则会在内部计算。默认为 None。

返回值:

与 A 形状相同的数组，包含在 A 的特征值上计算 func 的结果。

返回类型:

Array | tuple[Array, Array]

备注

JAX 实现的返回 dtype 可能与 scipy 的不同；具体来说，在数组值的所有虚部都接近于零的情况下，SciPy 函数可能会返回一个实值数组，而 JAX 实现将返回一个复值数组。

示例

应用任意矩阵函数

>>> A = jnp.array([[1., 2.], [3., 4.]])
>>> def func(x):
...   return jnp.sin(x) + 2 * jnp.cos(x)
>>> jax.scipy.linalg.funm(A, func)  
Array([[ 1.2452652 +0.j, -0.3701772 +0.j],
       [-0.55526584+0.j,  0.6899995 +0.j]], dtype=complex64)

比较计算矩阵指数的两种方法

>>> expA_1 = jax.scipy.linalg.funm(A, jnp.exp)
>>> expA_2 = jax.scipy.linalg.expm(A)
>>> jnp.allclose(expA_1, expA_2, rtol=1E-4)
Array(True, dtype=bool)