jax.scipy.linalg.funm

内容

jax.scipy.linalg.funm#

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

评估矩阵值函数

JAX 实现 scipy.linalg.funm().

参数:
  • A (ArrayLike) – 形状为 (N, N) 的数组,将为其计算函数。

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

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

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

返回:

A 形状相同的数组,包含在 A 的特征值上评估的 func 的结果。

返回类型:

Array | tuple[Array, Array]

笔记

JAX 实现的返回值类型可能与 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)