Preconditioners: custom callable inverse apply

Custom callables are for experimentation and interop. The contract is explicit: the callable or object with solve(x) must apply an approximation to A^{-1} @ x. mlx-sparse does not treat an arbitrary sparse matrix passed as M as something to invert. Custom callables use the host GMRES fallback and are therefore slower than native preconditioners.

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 user-supplied inverse diagonal

This example wraps a function with aspreconditioner. The shape is validated against A, and finite outputs are checked before the solver accepts the result.

A_sp = scipy.sparse.csr_array(
    np.array(
        [
            [6.0, -1.0, 0.0],
            [0.5, 5.0, -1.0],
            [0.0, 1.0, 4.0],
        ],
        dtype=np.float32,
    )
)
A = ms.from_scipy(A_sp)
b = mx.array([1.0, -2.0, 0.5], dtype=mx.float32)
inv_diag = mx.array(1.0 / A_sp.diagonal().astype(np.float32), dtype=mx.float32)

def inverse_apply(rhs):
    return inv_diag * rhs

M = preconditioners.aspreconditioner(inverse_apply, A)
x, info = linalg.gmres(A, b, M=M, rtol=1e-6, restart=3, maxiter=24, return_info=True)
print(M.kind, M.setup_device, M.apply_device)
print("GMRES status/iterations/residual:", info.status, info.iterations, f"{info.residual_norm:.3e}")

callable python_host python_host
GMRES status/iterations/residual: 0 3 9.127e-08

Production guidance

Use a native preconditioner when one matches the matrix class. Use custom callables for prototypes, external solver interop, or one-off diagnostics where Python inverse-apply cost is acceptable.