jax.scipy.sparse.linalg.cg#
- jax.scipy.sparse.linalg.cg(A, b, x0=None, *, tol=1e-05, atol=0.0, maxiter=None, M=None)[源代码]#
使用共轭梯度迭代法求解
Ax = b。JAX 的
cg的数值应该与 SciPy 的cg完全匹配（达到数值精度），但请注意，接口略有不同：您需要将线性算子A作为函数提供，而不是稀疏矩阵或LinearOperator。cg的导数通过隐式微分和另一个cg求解来实现，而不是通过微分通过求解器。 只有当两个求解都收敛时，它们才是准确的。- 参数:
A (ndarray, function, or matmul-compatible object) – 2D 数组或函数，用于计算线性映射（矩阵向量积）
Ax，当像A(x)或A @ x这样调用时。A必须表示 Hermitian 正定矩阵，并且必须返回与参数具有相同结构和形状的数组。b (数组 或 数组树) – 表示单个向量的线性系统的右侧。 可以存储为数组或具有任何形状的数组的 Python 容器。
x0 (array or tree of arrays) – 解的起始猜测。 必须与
b具有相同的结构。tol (float, optional) – 收敛容差，
norm(residual) <= max(tol*norm(b), atol)。 我们没有实现 SciPy 的“传统”行为，因此 JAX 的容差将与 SciPy 不同，除非您显式地将atol传递给 SciPy 的cg。atol (float, optional) – 收敛容差，
norm(residual) <= max(tol*norm(b), atol)。 我们没有实现 SciPy 的“传统”行为，因此 JAX 的容差将与 SciPy 不同，除非您显式地将atol传递给 SciPy 的cg。maxiter (integer) – 最大迭代次数。 即使未达到指定的容差，迭代也会在 maxiter 步后停止。
M (ndarray, function, or matmul-compatible object) – A 的预处理器。 预处理器应该近似 A 的逆。有效的预处理会显着提高收敛速度，这意味着需要更少的迭代才能达到给定的误差容限。
- 返回:
x (数组或数组树) – 收敛的解。 与
b具有相同的结构。info (None) – 收敛信息的占位符。 未来，JAX 将报告未实现收敛时的迭代次数，就像 SciPy 一样。