Preconditioners: explicit diagonal

preconditioners.diagonal(...) is for cases where you already have a safe inverse diagonal or a safe diagonal to invert. It is lower level than Jacobi: you provide the diagonal data directly, and the object applies an approximate inverse with native elementwise kernels for vector and matrix right-hand sides.

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 diagonal system

On a diagonal matrix, an exact inverse diagonal is an exact preconditioner. This is a useful smoke test because the preconditioned residual should collapse quickly and rank-2 application should simply scale each row of the RHS block.

diag = np.geomspace(1.0, 1.0e6, 24).astype(np.float32)
A = ms.diags(mx.array(diag, dtype=mx.float32), shape=(diag.size, diag.size))
b = mx.linspace(1.0, 2.0, diag.size)

M = preconditioners.diagonal(mx.array(1.0 / diag, dtype=mx.float32), inverse=True)
x, info = linalg.cg(A, b, M=M, rtol=1e-6, maxiter=16, return_info=True)

rhs_block = mx.stack([b, 2.0 * b], axis=1)
scaled_block = M(rhs_block)

print(M.kind, M.shape, M.apply_device)
print("CG status/iterations/residual:", info.status, info.iterations, f"{info.residual_norm:.3e}")
print("first scaled rows:", np.asarray(scaled_block[:3]).round(6))

diagonal (24, 24) native_cpu_or_metal
CG status/iterations/residual: 0 1 2.920e-07
first scaled rows: [[1.     2.    ]
 [0.5723 1.1446]
 [0.3269 0.6539]]

Safety notes

The constructor validates rank, shape, zero/near-zero entries, finite outputs, and rank-1/rank-2 RHS behavior. Passing inverse=False asks mlx-sparse to invert the provided diagonal explicitly, so zeros are rejected unless a documented zero policy is chosen.