Preconditioners: Jacobi

Jacobi builds an inverse diagonal from the summed sparse diagonal of A. It is the cheapest first-line preconditioner for scaled SPD systems and many diagonal dominant general systems. Setup uses existing sparse diagonal extraction; application is native elementwise work on the selected MLX device.

import mlx.core as mx
import numpy as np
import scipy.sparse
import mlx_sparse as ms
from mlx_sparse import linalg
from mlx_sparse.linalg import preconditioners

# Use CPU execution throughout.
ms.use_cpu(require_available=False)
np.set_printoptions(precision=4, suppress=True)

Scaled Poisson PCG

The matrix is D @ T @ D, where T is a tridiagonal SPD operator and D spans several orders of magnitude. Jacobi removes much of the scaling imbalance while keeping setup cost tiny.

n = 96
T = scipy.sparse.diags(
    [-np.ones(n - 1), 2.5 * np.ones(n), -np.ones(n - 1)],
    [-1, 0, 1],
    format="csr",
    dtype=np.float32,
)
scale = np.geomspace(1.0e-3, 1.0e3, n).astype(np.float32)
D = scipy.sparse.diags(scale, format="csr", dtype=np.float32)
A_sp = (D @ T @ D).astype(np.float32).tocsr()
A = ms.from_scipy(A_sp)
b = A @ mx.sin(mx.linspace(0.0, np.pi, n))

x0, info0 = linalg.cg(A, b, rtol=1e-5, maxiter=512, return_info=True)
M = preconditioners.jacobi(A, check=True)
xj, infoj = linalg.cg(A, b, M=M, rtol=1e-5, maxiter=512, return_info=True)

print("Jacobi nnz/fill:", M.nnz, M.nnz / A.nnz)
print("none  :", info0.status, info0.iterations, f"{info0.residual_norm:.3e}")
print("jacobi:", infoj.status, infoj.iterations, f"{infoj.residual_norm:.3e}")

Jacobi nnz/fill: 96 0.3356643356643357
none  : 512 512 9.908e+00
jacobi: 0 15 3.400e-01

GMRES use

GMRES uses left preconditioning for Jacobi: it builds the Krylov basis for M^{-1} A but still tests convergence on the true residual b - A @ x. This matches the standard SciPy/PETSc convention and avoids false success based only on a preconditioned residual.