Preconditioners: identity baseline#

The identity preconditioner is the release sanity check. It gives the solver the same mathematical problem as M=None, but exercises the public preconditioner protocol: shape, dtype, metadata, inverse-apply validation, and solver dispatch. Use it when you need a controlled baseline before introducing Jacobi, ILU(0), IC(0), Chebyshev, or exact factors.

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)

A small SPD system#

The matrix below is a 1-D Poisson operator with a mass shift. It is symmetric positive-definite, so CG is the natural solver. We solve the same system with no preconditioner and with identity(A), then compare the structured diagnostics.

n = 32
main = 2.25 * np.ones(n, dtype=np.float32)
off = -1.0 * np.ones(n - 1, dtype=np.float32)
A_sp = scipy.sparse.diags([off, main, off], [-1, 0, 1], format="csr", dtype=np.float32)
A = ms.from_scipy(A_sp)
x_true = mx.sin(mx.linspace(0.0, np.pi, n))
b = A @ x_true

x_none, info_none = linalg.cg(A, b, rtol=1e-6, maxiter=128, return_info=True)
M = preconditioners.identity(A)
x_id, info_id = linalg.cg(A, b, M=M, rtol=1e-6, maxiter=128, return_info=True)

print(M.kind, M.shape, M.dtype, M.setup_device, M.apply_device)
print("none    ", info_none.status, info_none.iterations, f"{info_none.residual_norm:.3e}")
print("identity", info_id.status, info_id.iterations, f"{info_id.residual_norm:.3e}")
print("solutions close:", bool(mx.allclose(x_none, x_id, rtol=1e-6, atol=1e-6)))
identity (32, 32) mlx.core.float32 none none
none     0 16 3.057e-09
identity 0 16 3.057e-09
solutions close: True

What this validates#

  • identity(A) is a protocol object, not a sparse matrix inversion request.

  • cg(..., M=identity(A)) preserves the default (x, info) API unless return_info=True is requested.

  • A real release benchmark should report identity beside the unpreconditioned solver so preconditioner overhead and dispatch behavior are visible.