Sparse linalg API#

The linalg namespace contains sparse-native solvers and factorizations. Public functions accept CSRArray, COOArray, and CSCArray inputs, dense arrays are not silently converted because dense linear algebra belongs in mlx.linalg. COO and CSC inputs are normalized once at solver entry. The native sparse kernels use CSR, Accelerate-enabled direct solves normalize supported real inputs to canonical CSC before calling Apple’s framework.

See also

See Solvers for the solver support matrix, including Full CPU + GPU, CPU only, GPU only, Partial, and Accelerate coverage labels.

Iterative solvers#

mlx_sparse.linalg.cg(A, b, *, x0=None, rtol=1e-05, atol=0.0, maxiter=None, M=None, callback=None, return_info=False)[source]#

Solve a sparse SPD linear system with the conjugate gradient method.

Conjugate gradients (CG) is an iterative Krylov solver for symmetric positive-definite (SPD) systems A @ x = b. Each iteration requires one sparse matrix-vector product dispatched to a Metal kernel via the native CSR engine. CG converges in at most n steps in exact arithmetic and typically far fewer for well-conditioned problems.

GPU note:

When GPU execution is selected, the native CG iteration runs in a Metal kernel. Sparse matrix-vector products, vector updates, dot products, and residual checks stay on the GPU. Python argument validation and conversion of the returned info flag to a Python integer happen on the host.

Parameters:

  • A – Coefficient matrix. Must be a CSRArray, COOArray, or CSCArray that is symmetric positive-definite. A sparse-backed LinearOperator is also accepted. A fully matrix-free LinearOperator uses a documented host fallback loop.

  • b – Right-hand side vector of shape (n,). Must be a rank-1 mlx.core.array or anything convertible to one.

  • x0 – Initial guess of shape (n,). Defaults to the zero vector when None.

  • rtol (float) – Relative tolerance. The solver stops when ||r_k|| <= max(atol, rtol * ||b||). Defaults to 1e-5.

  • atol (float) – Absolute tolerance floor. Useful when b is near zero. Defaults to 0.0.

  • maxiter (int | None) – Maximum number of iterations. Defaults to 10 * n when None.

  • M – Optional native-backed preconditioner. None uses the existing unpreconditioned native CG path. preconditioners.identity also uses that path. preconditioners.diagonal and preconditioners.jacobi dispatch to native Jacobi-preconditioned CG. preconditioners.ichol0 dispatches to a native CPU IC(0)-preconditioned CG loop whose preconditioner application uses the stored incomplete Cholesky factors. preconditioners.chebyshev dispatches to a native Chebyshev-polynomial PCG loop whose preconditioner application uses only sparse matrix-vector products and vector updates. Fully matrix-free LinearOperator inputs accept arbitrary inverse-apply objects and callables through a slower host fallback.

  • callback – Optional callable invoked once after the native solve completes. Native CPU/Metal Krylov loops do not call Python inside each iteration; using a callback synchronizes only the final solution.

  • return_info (bool) – If True, return a structured diagnostic object instead of the integer info flag. The default remains False.

Returns:

A tuple (x, info) where x is the approximate solution array of shape (n,) and info is an integer convergence flag by default. With return_info=True, info is a SolverInfo diagnostic object. info == 0 means the solver converged to the requested tolerance. info > 0 is the iteration count at which the solver stopped without converging. info < 0 indicates numerical breakdown, such as a non-positive preconditioned residual product, non-finite scalar, or scale-aware near-zero denominator.

Raises:

  • TypeError – If A is a dense mlx.core.array or an unsupported type, or if M is not a native-backed identity, diagonal, or Jacobi preconditioner.

  • ValueError – If b is not rank-1 or its length does not match A.shape[0].

Example

Build a 2-D Poisson Laplacian and solve it with CG:

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

n = 16
L = scipy.sparse.diags([-1, 4, -1], [-1, 0, 1],
                       shape=(n, n), format='csr').astype(np.float32)
A = ms.csr_array(
    (mx.array(L.data), mx.array(L.indices), mx.array(L.indptr)),
    shape=L.shape, canonical=True,
)
b = mx.ones((n,), dtype=mx.float32)
x, info = linalg.cg(A, b, rtol=1e-6)
print(info)  # 0 means converged
mlx_sparse.linalg.gmres(A, b, *, x0=None, rtol=1e-05, atol=0.0, restart=None, maxiter=None, M=None, callback=None, callback_type='x', return_info=False)[source]#

Solve a sparse linear system with restarted GMRES.

Generalised Minimum RESidual (GMRES) is an iterative Krylov solver for any non-singular square system A @ x = b. Unlike CG it does not require A to be symmetric or positive-definite, making it the default choice for non-symmetric PDEs and general sparse systems.

Each restart cycle builds an Arnoldi basis of dimension restart using one sparse matrix-vector product per step, then minimises the residual over that Krylov subspace. Memory scales as O(restart * n).

GPU note:

When GPU execution is selected, Arnoldi basis construction uses the native Arnoldi kernel. The restart loop, residual setup, small least-squares solve, solution update, and convergence bookkeeping run on the CPU. The host also copies the Arnoldi basis and Hessenberg coefficients back before updating the solution.

Parameters:

  • A – Coefficient matrix. Must be a CSRArray, COOArray, or CSCArray. A sparse-backed LinearOperator is also accepted. A fully matrix-free LinearOperator uses a documented host fallback loop. Need not be symmetric.

  • b – Right-hand side vector of shape (n,).

  • x0 – Initial guess of shape (n,). Defaults to the zero vector.

  • rtol (float) – Relative tolerance for the residual stopping criterion. Defaults to 1e-5.

  • atol (float) – Absolute tolerance floor. Defaults to 0.0.

  • restart (int | None) – Size of the Krylov subspace built before each restart. Larger values accelerate convergence at the cost of more memory. Defaults to min(20, n).

  • maxiter (int | None) – Maximum total number of matrix-vector products across all restart cycles. Defaults to 10 * n.

  • M – Optional inverse-apply preconditioner. None uses the existing native unpreconditioned GMRES path. preconditioners.identity is treated equivalently. preconditioners.diagonal and preconditioners.jacobi and preconditioners.ilu0 dispatch to native left-preconditioned GMRES. Custom callables and inverse-apply objects use a host fallback loop. Fully matrix-free LinearOperator inputs also accept arbitrary inverse-apply objects and callables through the same slower host loop. All preconditioned paths build Krylov vectors for M^{-1} A while convergence is tested against the true residual b - A @ x.

  • callback – Optional callable invoked once after the native solve completes. Native CPU/Metal solver loops do not call Python inside each iteration. callback_type="x" receives the final solution; "pr_norm" and "legacy" receive the final reported residual norm. The payload names mirror SciPy’s callback vocabulary, but native mlx-sparse paths intentionally do not call Python at every restart or inner iteration.

  • callback_type (str) – One of "x", "pr_norm", or "legacy". "legacy" does not change maxiter accounting.

  • return_info (bool) – If True, return a structured diagnostic object instead of the integer info flag. The default remains False.

Returns:

A tuple (x, info) where x is the approximate solution of shape (n,) and info is an integer convergence flag by default. With return_info=True, info is a SolverInfo diagnostic object. info == 0 means the solver converged. info > 0 is the iteration count at termination without convergence.

Raises:

  • TypeError – If A is a dense array or an unsupported type.

  • ValueError – If b is not rank-1, its length does not match A.shape[0], or restart is not positive.

mlx_sparse.linalg.minres(A, b, *, x0=None, rtol=1e-05, atol=0.0, shift=0.0, maxiter=None, M=None, check_preconditioner=True, callback=None, return_info=False)[source]#

Solve a sparse symmetric linear system with MINRES.

MINimum RESidual (MINRES) is an iterative Krylov solver for symmetric systems A @ x = b where A may be symmetric indefinite (having both positive and negative eigenvalues). It minimises the 2-norm of the residual at every step using a Paige-Saunders Lanczos recurrence and requires only a constant number of vectors in memory, making it more memory-efficient than GMRES for large symmetric indefinite problems.

GPU note:

When GPU execution is selected, the unpreconditioned and diagonal/Jacobi-preconditioned recurrences dispatch to native Metal kernels. The CPU path uses the matching native C++ recurrence. No Python callback is invoked inside the Krylov loop.

Parameters:

  • A – Coefficient matrix. Must be a CSRArray, COOArray, or CSCArray that is real and symmetric.

  • b – Right-hand side vector of shape (n,).

  • x0 – Initial guess of shape (n,). Defaults to the zero vector.

  • rtol (float) – Relative tolerance for the residual stopping criterion. Defaults to 1e-5.

  • atol (float) – Absolute tolerance floor. Defaults to 0.0.

  • shift (float) – Optional scalar shift. The solver applies MINRES to (A - shift * I) @ x = b, matching SciPy’s convention.

  • maxiter (int | None) – Maximum number of iterations. Defaults to 10 * n.

  • M – Optional symmetric positive-definite inverse-apply preconditioner. None and preconditioners.identity use the unpreconditioned path. preconditioners.diagonal and preconditioners.jacobi dispatch to native diagonal-preconditioned MINRES. Other preconditioner kinds are rejected until they have native SPD MINRES kernels.

  • check_preconditioner (bool) – When True (default), diagonal/Jacobi preconditioners must have finite strictly positive inverse diagonal entries before entering native MINRES. Setting this to False disables the Python-side SPD validation but the native solver still reports numerical breakdown for invalid preconditioners.

  • callback – Optional callable invoked once after the native solve completes. Native CPU/Metal Krylov loops do not call Python inside each iteration; using a callback synchronizes only the final solution.

  • return_info (bool) – If True, return a structured diagnostic object instead of the integer info flag. The default remains False.

Returns:

A tuple (x, info) where x is the approximate solution of shape (n,) and info is an integer convergence flag by default. With return_info=True, info is a SolverInfo diagnostic object. info == 0 means the solver converged to the requested tolerance. info > 0 is the iteration count at which the solver stopped without converging.

Raises:

  • TypeError – If A is a dense array or an unsupported type.

  • ValueError – If b is not rank-1, its length does not match A.shape[0], or a checked diagonal preconditioner is not SPD.

Note

MINRES requires A to be symmetric but not positive-definite. Preconditioned MINRES additionally requires an SPD preconditioner. For SPD systems cg() is typically faster. For non-symmetric systems use gmres().

Preconditioners#

Native-backed sparse solver preconditioners.

The Python objects in this module are containers and dispatch helpers. Application and Krylov iteration dispatch to native mlx-sparse primitives rather than Python solver loops. Constructors may use existing sparse native kernels and MLX scalar array expressions to build immutable setup data.

class mlx_sparse.linalg.preconditioners.DiagonalPreconditioner(inverse_diagonal, shape, kind='diagonal', is_symmetric=True, is_positive_definite=False)[source]#

Bases: object

Explicit diagonal inverse-apply preconditioner.

The stored vector is the inverse diagonal that should multiply each row of a right-hand side. Application dispatches to the native diagonal_preconditioner_apply primitive, including rank-2 RHS support.

Stored fields include the finite float32 inverse_diagonal vector, the compatible square shape, the stable kind string, and symmetry/positive-definiteness metadata.

Parameters:

  • inverse_diagonal (mlx.core.array)

  • shape (tuple[int, int])

  • kind (str)

  • is_symmetric (bool)

  • is_positive_definite (bool)

inverse_diagonal: mlx.core.array#
shape: tuple[int, int]#
kind: str#
is_symmetric: bool#
is_positive_definite: bool#
__post_init__()[source]#

Validate shape and inverse diagonal storage.

Return type:

None

property dtype#

Value dtype of inverse_diagonal.

property nnz: int#

Number of stored inverse diagonal entries.

property setup_device: str#

Device category used for setup validation.

property apply_device: str#

Device category used for inverse application.

property setup_info: Mapping[str, object]#

Structured metadata describing setup choices and assumptions.

solve(x)[source]#

Apply the diagonal inverse to a vector or dense RHS matrix.

Parameters:

x – Right-hand side with shape (n,) or (n, nrhs).

Returns:

Native-applied inverse_diagonal[:, None] * x for matrix RHS, or inverse_diagonal * x for vector RHS.

Return type:

mlx.core.array

matvec(x)[source]#

Alias for solve() for SciPy-style inverse-operator use.

Return type:

mlx.core.array

__call__(x)[source]#

Alias for solve().

Return type:

mlx.core.array

class mlx_sparse.linalg.preconditioners.CallablePreconditioner(apply, shape, dtype=mlx.core.float32, kind='callable', is_symmetric=False, is_positive_definite=False)[source]#

Bases: object

Python host inverse-apply preconditioner wrapper.

CallablePreconditioner is the explicit normalization layer for custom inverse-apply objects. The wrapped callable receives a rank-1 or rank-2 float32 MLX array and must return the same shape containing M^{-1} @ x. Solver integrations may use this wrapper only on documented host fallback paths because each application crosses through Python.

Stored fields include the callable apply object, compatible square shape, stable kind metadata, conservative symmetry/positive definiteness flags, and structured setup information.

Parameters:
apply: object#
shape: tuple[int, int]#
dtype: object#
kind: str#
is_symmetric: bool#
is_positive_definite: bool#
__post_init__()[source]#

Validate the callable contract metadata.

Return type:

None

property nnz: int#

Unknown effective storage count for custom callables.

property setup_device: str#

Device category used during setup.

property apply_device: str#

Device category used during inverse application.

property setup_info: Mapping[str, object]#

Structured metadata describing the callable contract.

solve(x)[source]#

Apply the wrapped inverse callable and validate its output.

Parameters:

x – Right-hand side with shape (n,) or (n, nrhs).

Returns:

Finite float32 output with the exact same shape as x.

Return type:

mlx.core.array

__call__(x)[source]#

Alias for solve().

Return type:

mlx.core.array

class mlx_sparse.linalg.preconditioners.ChebyshevPreconditioner(A, degree=2, lambda_min=0.0, lambda_max=0.0, estimate=True, spectral_info=<factory>, kind='chebyshev', is_symmetric=True, is_positive_definite=True)[source]#

Bases: object

Polynomial inverse preconditioner for SPD sparse matrices.

ChebyshevPreconditioner stores a canonical float32 CSR operator and applies a fixed-degree first-kind Chebyshev semi-iteration with zero initial guess. Application uses only sparse matrix-vector/matrix products and vector updates, so it follows the selected native CPU or Metal device path without triangular solves or dense factorization.

Stored fields include the CSR operator, polynomial degree, validated positive spectral interval [lambda_min, lambda_max], whether setup used native spectral estimate data, and detailed spectral_info metadata.

Parameters:
A: CSRArray#
degree: int#
lambda_min: float#
lambda_max: float#
estimate: bool#
spectral_info: Mapping[str, object]#
kind: str#
is_symmetric: bool#
is_positive_definite: bool#
__post_init__()[source]#

Validate polynomial metadata and CSR storage.

Return type:

None

property shape: tuple[int, int]#

Shape of the preconditioned square operator.

property dtype#

Value dtype used by the stored operator.

property nnz: int#

Number of sparse operator entries used during polynomial apply.

property setup_device: str#

Device category used during spectral setup.

property apply_device: str#

Device category used during inverse application.

property setup_info: Mapping[str, object]#

Structured metadata describing Chebyshev setup choices.

solve(x)[source]#

Apply the Chebyshev polynomial inverse to a vector or matrix RHS.

Parameters:

x – Right-hand side with shape (n,) or (n, nrhs).

Returns:

Native Chebyshev semi-iteration result with the same shape as x.

Return type:

mlx.core.array

matvec(x)[source]#

Alias for solve() for inverse-operator composition.

Return type:

mlx.core.array

__call__(x)[source]#

Alias for solve().

Return type:

mlx.core.array

class mlx_sparse.linalg.preconditioners.ExactFactorPreconditioner(solver, shape, method, backend, kind='exact', is_symmetric=False, is_positive_definite=False, factor_nnz=-1, native_apply_kind=None, native_factorization=None)[source]#

Bases: object

Exact inverse-apply preconditioner backed by a sparse factorization.

This wrapper composes existing direct sparse solve objects with the iterative M protocol. It does not refactorize on application: setup is completed before construction. solve() uses explicit native LU/Cholesky apply bindings when factors are available, uses the guarded Accelerate solver for real Accelerate factorized objects, and otherwise delegates to the stored reusable solve object.

Stored fields include the reusable solver, compatible square shape, factorization method, implementation backend, and conservative symmetry/positive-definiteness metadata. Accelerate-backed FactorizedSolve instances keep their Accelerate CPU apply boundary; native explicit factors keep their native CPU/Metal triangular-solve apply boundary.

Parameters:
solver: object#
shape: tuple[int, int]#
method: str#
backend: str#
kind: str#
is_symmetric: bool#
is_positive_definite: bool#
factor_nnz: int#
native_apply_kind: str | None#
native_factorization: object | None#
__post_init__()[source]#

Validate the wrapped exact factorization metadata.

Return type:

None

property dtype#

Value dtype used by current sparse direct solve backends.

property nnz: int#

Stored factor nonzero count, or -1 for opaque factors.

property setup_device: str#

Device category used during factorization setup.

property apply_device: str#

Device category used during inverse application.

property setup_info: Mapping[str, object]#

Structured metadata describing the exact factorization wrapper.

solve(x)[source]#

Apply the exact factorized solve to a vector or dense RHS matrix.

Parameters:

x – Right-hand side with shape (n,) or (n, nrhs).

Returns:

Finite float32 solution with the same shape as x.

Return type:

mlx.core.array

matvec(x)[source]#

Alias for solve() for inverse-operator composition.

Return type:

mlx.core.array

__call__(x)[source]#

Alias for solve().

Return type:

mlx.core.array

class mlx_sparse.linalg.preconditioners.IC0Preconditioner(L, shift=0.0, check=True, kind='ichol0', is_symmetric=True, is_positive_definite=True)[source]#

Bases: object

Natural-order no-fill incomplete Cholesky preconditioner.

IC0Preconditioner stores an explicit lower CSR factor L produced by native CPU IC(0) setup. The factor uses natural ordering, preserves the symmetric lower sparsity pattern of the canonical CSR input, and introduces no fill. Application performs two native triangular solves, L y = x followed by L.T z = y.

Stored fields include the lower factor, explicit non-negative diagonal shift used during setup, strict positive-pivot check mode, and SPD metadata suitable for CG and MINRES-style solvers.

Parameters:
L: CSRArray#
shift: float#
check: bool#
kind: str#
is_symmetric: bool#
is_positive_definite: bool#
__post_init__()[source]#

Validate factor metadata and explicit shift policy.

Return type:

None

property shape: tuple[int, int]#

Shape of the preconditioned square operator.

property dtype#

Value dtype used by the stored factor.

property nnz_L: int#

Stored nonzero count in the lower factor.

property nnz: int#

Total stored factor entries.

property setup_device: str#

Device category used during IC(0) setup.

property apply_device: str#

Device category used during inverse application.

property setup_info: Mapping[str, object]#

Structured metadata describing IC(0) setup choices.

solve(x)[source]#

Apply the IC(0) inverse approximation to a vector or matrix RHS.

Parameters:

x – Right-hand side with shape (n,) or (n, nrhs).

Returns:

Native triangular-solve result L.T^{-1} L^{-1} x with the same rank and leading dimension as x.

Return type:

mlx.core.array

matvec(x)[source]#

Alias for solve() for inverse-operator composition.

Return type:

mlx.core.array

__call__(x)[source]#

Alias for solve().

Return type:

mlx.core.array

class mlx_sparse.linalg.preconditioners.ILU0Preconditioner(L, U, shift=0.0, check=True, reuse_analysis=False, kind='ilu0', is_symmetric=False, is_positive_definite=False)[source]#

Bases: object

Natural-order no-fill incomplete LU preconditioner.

ILU0Preconditioner stores explicit CSR factors L and U from a native CPU ILU(0) setup. L has a unit diagonal and the factors preserve the lower/upper sparsity pattern of the canonical CSR input without fill-reducing ordering or pivoting. Application performs the standard forward and backward triangular solves L y = x and U z = y using native CSR triangular-solve kernels.

Stored fields include the lower/upper CSR factors, explicit diagonal shift used during setup, validation check mode, optional triangular-analysis reuse flag, and conservative nonsymmetric metadata.

Parameters:
L: CSRArray#
U: CSRArray#
shift: float#
check: bool#
reuse_analysis: bool#
kind: str#
is_symmetric: bool#
is_positive_definite: bool#
__post_init__()[source]#

Validate factor metadata and optionally cache triangular analysis.

Return type:

None

property shape: tuple[int, int]#

Shape of the preconditioned square operator.

property dtype#

Value dtype used by the stored factors.

property nnz_L: int#

Stored nonzero count in the unit lower factor.

property nnz_U: int#

Stored nonzero count in the upper factor.

property nnz: int#

Total stored factor entries.

property setup_device: str#

Device category used during ILU(0) setup.

property apply_device: str#

Device category used during inverse application.

property setup_info: Mapping[str, object]#

Structured metadata describing ILU(0) setup choices.

solve(x)[source]#

Apply the ILU(0) inverse approximation to a vector or matrix RHS.

Parameters:

x – Right-hand side with shape (n,) or (n, nrhs).

Returns:

Native triangular-solve result U^{-1} L^{-1} x with the same rank and leading dimension as x.

Return type:

mlx.core.array

matvec(x)[source]#

Alias for solve() for inverse-operator composition.

Return type:

mlx.core.array

__call__(x)[source]#

Alias for solve().

Return type:

mlx.core.array

class mlx_sparse.linalg.preconditioners.IdentityPreconditioner(shape, dtype=mlx.core.float32, kind='identity', is_symmetric=True, is_positive_definite=True)[source]#

Bases: object

No-op inverse-apply preconditioner.

IdentityPreconditioner is useful as an explicit baseline and as the normalized representation of M=None inside solver plumbing.

Stored fields include the compatible square shape, the native value dtype (currently mlx.core.float32), the stable kind string, and symmetry/positive-definiteness metadata.

Parameters:
shape: tuple[int, int]#
dtype: object#
kind: str#
is_symmetric: bool#
is_positive_definite: bool#
__post_init__()[source]#

Normalize and validate the stored square shape.

Return type:

None

property nnz: int#

Number of effective diagonal entries.

property setup_device: str#

Device used during setup.

property apply_device: str#

Device used during inverse application.

property setup_info: Mapping[str, object]#

Structured metadata describing setup choices and assumptions.

solve(x)[source]#

Return x after validating rank, shape, dtype, and finiteness.

Parameters:

x – Right-hand side vector (n,) or matrix (n, nrhs).

Returns:

x as a finite float32 MLX array.

Return type:

mlx.core.array

__call__(x)[source]#

Alias for solve().

Return type:

mlx.core.array

class mlx_sparse.linalg.preconditioners.JacobiPreconditioner(inverse_diagonal, shape, kind='jacobi', is_symmetric=True, is_positive_definite=False, omega=1.0, shift=0.0, zero_policy='raise', zero_atol=0.0, checked=False, positive_diagonal=None)[source]#

Bases: DiagonalPreconditioner

Jacobi preconditioner built from a sparse matrix diagonal.

JacobiPreconditioner is a specialized diagonal inverse-apply object that records the setup parameters used by jacobi(). Passing it to mlx_sparse.linalg.cg() dispatches to the native Jacobi-PCG primitive.

Stored fields include omega, shift, zero_policy, zero_atol, whether validation was checked, and positive_diagonal when the cheap positive shifted-diagonal check was requested.

Parameters:
omega: float#
shift: float#
zero_policy: str#
zero_atol: float#
checked: bool#
positive_diagonal: bool | None#
property setup_device: str#

Device category used for Jacobi setup.

property setup_info: Mapping[str, object]#

Structured metadata describing Jacobi setup choices.

class mlx_sparse.linalg.preconditioners.Preconditioner(*args, **kwargs)[source]#

Bases: Protocol

Protocol for objects that apply an approximate inverse.

Solver-facing preconditioners represent the operation M^{-1} @ x, not the matrix M itself. Implementations expose a square shape, value dtype, a stable kind identifier, symmetry and positive-definiteness metadata, setup/apply device descriptors, effective storage nnz, and a structured setup_info mapping. The required solve() method and __call__ alias must accept rank-1 vector RHS and rank-2 dense RHS matrices without mutating the matrix used during setup.

shape: tuple[int, int]#
dtype: object#
kind: str#
is_symmetric: bool#
is_positive_definite: bool#
setup_device: str#
apply_device: str#
nnz: int#
setup_info: Mapping[str, object]#
solve(x)[source]#

Apply the preconditioner solve to x.

Return type:

mlx.core.array

__call__(x)[source]#

Alias for solve().

Return type:

mlx.core.array

mlx_sparse.linalg.preconditioners.aspreconditioner(M, A=None, *, assume_inverse=True)[source]#

Normalize supported preconditioner-like objects.

Parameters:

  • M – None, an existing preconditioner, an object with solve(x), or a callable. Sparse matrices are rejected because they do not explicitly define an inverse-apply contract.

  • A – Optional reference matrix or shape used to validate compatibility.

  • assume_inverse (bool) – Must be True for callables and custom objects, documenting that their output is already an inverse/preconditioner application.

Returns:

A supported preconditioner object.

Raises:

  • ValueError – If A is required or the preconditioner shape mismatches.

  • TypeError – If M is a sparse matrix or unsupported object.

Return type:

Preconditioner

mlx_sparse.linalg.preconditioners.chebyshev(A, *, degree=2, lambda_min=None, lambda_max=None, estimate=True)[source]#

Create a GPU-friendly Chebyshev polynomial preconditioner.

The setup normalizes A to canonical CSR, promotes real low-precision values to float32, and computes native Gershgorin spectral bounds plus optional native Lanczos Ritz estimates. The resulting preconditioner applies a fixed-degree first-kind Chebyshev semi-iteration with zero initial guess, using only sparse matrix products and vector updates.

Parameters:

  • A – CSRArray, COOArray, CSCArray, or sparse-backed LinearOperator describing a real SPD square matrix.

  • degree (int) – Positive polynomial degree. Defaults to 2.

  • lambda_min (float | None) – Optional positive lower spectral bound. If omitted, setup uses a positive Gershgorin lower bound when available, otherwise a conservative Lanczos-derived lower estimate when estimate=True.

  • lambda_max (float | None) – Optional upper spectral bound. If omitted, setup uses the Gershgorin upper bound when available, with Lanczos metadata recorded for diagnostics.

  • estimate (bool) – Whether native Lanczos Ritz estimates should be computed as a fallback/refinement for the spectral interval. Defaults to True.

Returns:

A ChebyshevPreconditioner with native CPU/Metal apply support.

Raises:

ValueError – If no valid positive interval can be established.

Return type:

ChebyshevPreconditioner

mlx_sparse.linalg.preconditioners.diagonal(inv_diag_or_diag, *, inverse=False, shape=None, dtype=None, zero_atol=0.0)[source]#

Create an explicit diagonal inverse-apply preconditioner.

Parameters:

  • inv_diag_or_diag – Rank-1 diagonal values. Interpreted as a diagonal by default, or as an inverse diagonal when inverse=True.

  • inverse (bool) – If True, use inv_diag_or_diag directly as the inverse diagonal. If False, validate and invert it.

  • shape – Optional square shape. Defaults to (n, n) where n is the vector length.

  • dtype – Optional dtype. The current native preconditioner path accepts only None or mlx.core.float32.

  • zero_atol (float) – Absolute threshold used when rejecting zero diagonal entries before inversion.

Returns:

A DiagonalPreconditioner with finite float32 inverse diagonal storage.

Return type:

DiagonalPreconditioner

mlx_sparse.linalg.preconditioners.exact(A, *, method='auto')[source]#

Factorize A once and return an exact inverse-apply preconditioner.

exact is a convenience wrapper around mlx_sparse.linalg.factorized(). It is intended as a correctness baseline, diagnostic tool, and composition point for existing direct solvers rather than a performance headline.

Parameters:
Returns:

An ExactFactorPreconditioner wrapping the reusable factorized solve object.

Return type:

ExactFactorPreconditioner

mlx_sparse.linalg.preconditioners.from_factorized(solver)[source]#

Wrap an existing sparse factorization as an exact preconditioner.

Parameters:

solver – A FactorizedSolve, SparseLU, or SparseCholesky instance.

Returns:

An ExactFactorPreconditioner whose inverse application uses a native exact-apply path when the factorization exposes one.

Raises:

  • TypeError – If solver is not one of the supported factorization objects.

  • ValueError – If the factorization does not represent a square operator.

Return type:

ExactFactorPreconditioner

mlx_sparse.linalg.preconditioners.identity(A_or_shape, *, dtype=None)[source]#

Create a no-op preconditioner for a square shape or sparse matrix.

Parameters:

  • A_or_shape – Square sparse matrix, (n, n) shape tuple, or integer dimension.

  • dtype – Optional dtype. The current native solver integration accepts only None or mlx.core.float32.

Returns:

An IdentityPreconditioner.

Return type:

IdentityPreconditioner

mlx_sparse.linalg.preconditioners.ichol0(A, *, shift=0.0, check=True)[source]#

Create a natural-order no-fill IC(0) preconditioner.

The setup normalizes A to canonical CSR, promotes real low-precision values to float32, and runs a native CPU incomplete Cholesky factorization with zero fill and natural ordering. The original sparse matrix is not mutated. shift is a non-negative scalar added only to existing diagonal entries before setup; a missing diagonal remains an error and no pivot is silently perturbed.

Parameters:

  • A – CSRArray, COOArray, CSCArray, or sparse-backed LinearOperator describing an SPD square matrix.

  • shift (float) – Explicit non-negative diagonal shift added before factorization. Defaults to 0.0.

  • check (bool) – When True (default), native setup enforces symmetry for explicitly stored mirrored entries and uses a scale-aware positive pivot guard. When False, non-positive and non-finite pivots are still rejected, but near-zero pivot and symmetry checks are relaxed.

Returns:

An IC0Preconditioner whose application uses native CSR triangular solves L and L.T.

Return type:

IC0Preconditioner

mlx_sparse.linalg.preconditioners.ilu0(A, *, shift=0.0, check=True, reuse_analysis=False)[source]#

Create a natural-order no-fill ILU(0) preconditioner.

The setup normalizes A to canonical CSR, promotes real low-precision values to float32, and runs a native CPU ILU(0) factorization with no fill-reducing ordering and no pivoting. The original sparse matrix is not mutated. shift is added only to existing diagonal entries before setup; a missing diagonal remains an error.

Parameters:

  • A – CSRArray, COOArray, CSCArray, or sparse-backed LinearOperator describing a square nonsingular matrix.

  • shift (float) – Explicit diagonal shift added before factorization.

  • check (bool) – When True (default), native setup uses a scale-aware near-zero pivot guard. When False, only exact zero and non-finite pivots are rejected during setup.

  • reuse_analysis (bool) – When True, cache triangular diagonal-position and level-schedule analysis for repeated explicit M(rhs) calls. The default is False because v0.0.4b1 triangular-analysis benchmarks showed this must be workload-measured before enabling.

Returns:

An ILU0Preconditioner whose application uses native CSR triangular solves.

Return type:

ILU0Preconditioner

mlx_sparse.linalg.preconditioners.jacobi(A, *, omega=1.0, shift=0.0, zero_policy='raise', zero_atol=0.0, check=False)[source]#

Create a Jacobi preconditioner from a sparse matrix diagonal.

The inverse diagonal is computed as omega / (diag(A) + shift) after normalizing A to canonical CSR so duplicate diagonal entries are summed. The input sparse matrix is never mutated.

Parameters:

  • A – CSRArray, COOArray, CSCArray, or sparse-backed LinearOperator.

  • omega (float) – Damping/weighting factor.

  • shift (float) – Explicit diagonal shift applied before inversion.

  • zero_policy (str) – "raise" rejects zero/near-zero shifted diagonals. "unit" replaces those entries with 1 before inversion.

  • zero_atol (float) – Absolute threshold used to identify near-zero shifted diagonal entries.

  • check (bool) – When True, require omega > 0 and a strictly positive shifted diagonal before any zero_policy replacement, then mark the preconditioner as positive definite.

Returns:

A JacobiPreconditioner suitable for native PCG.

Return type:

JacobiPreconditioner

Spectral routines#

mlx_sparse.linalg.lanczos(A, k, *, v0=None, reorthogonalize=True, return_basis=True)[source]#

Run the Lanczos iteration on a sparse symmetric matrix.

Builds a Krylov basis of dimension k for the symmetric matrix A using the Lanczos three-term recurrence. The basis vectors, together with the tridiagonal matrix they define, contain the information needed to approximate eigenvalues and eigenvectors via the Ritz pairs of A.

This is a low-level routine. For eigenvalue computation, prefer eigsh() which calls Lanczos internally and returns the final eigenpairs directly.

GPU note:

When GPU execution is selected, the Lanczos recurrence runs in a Metal kernel. Sparse matrix-vector products, orthogonalisation, tridiagonal coefficients, and basis writes stay on the GPU. Python argument validation and returned array handling happen on the host.

Parameters:

  • A – Symmetric sparse matrix. Must be a CSRArray, COOArray, or CSCArray. Float16 and bfloat16 inputs are promoted to float32.

  • k (int) – Number of Lanczos steps. Must satisfy 0 < k <= A.shape[0].

  • v0 – Optional starting vector of shape (n,). None uses the deterministic all-ones start vector.

  • reorthogonalize (bool) – Whether to apply full reorthogonalisation at each step to suppress numerical loss of orthogonality. Defaults to True.

  • return_basis (bool) – When True (the default), return the Lanczos basis matrix Q in addition to the tridiagonal coefficients.

Returns:

When return_basis=True, a tuple (alphas, betas, Q) where alphas is the diagonal of shape (k,), betas is the sub-diagonal of shape (k-1,) or (k,), and Q is the basis matrix of shape (n, k). When return_basis=False, returns (alphas, betas) without the basis.

Raises:

ValueError – If k is out of range.

mlx_sparse.linalg.eigsh(A, k=6, *, which='LM', v0=None, ncv=None, maxiter=None, tol=0.0, return_eigenvectors=True)[source]#

Compute a few eigenpairs of a sparse symmetric matrix.

Uses the native CSR Lanczos-based eigensolver to find the k eigenpairs of the real symmetric (Hermitian) sparse matrix A that match the criterion specified by which. Each Lanczos step dispatches a sparse matrix-vector product to the GPU via the native Metal kernel.

GPU note:

When GPU execution is selected, Lanczos tridiagonalisation uses the native Lanczos kernel. The small tridiagonal eigensolve, Ritz pair selection, and eigenvector back transformation run on the CPU after the basis and coefficients are copied back to host memory.

Parameters:

  • A – Real symmetric sparse matrix. Must be a CSRArray, COOArray, or CSCArray. Float16 and bfloat16 inputs are promoted to float32.

  • k (int) – Number of eigenpairs to compute. Must satisfy 0 < k < A.shape[0]. Defaults to 6.

  • which (str) –

    Which eigenpairs to return. Accepted values:

    • "LM": eigenvalues Largest in Magnitude (default)

    • "SM": eigenvalues Smallest in Magnitude

    • "LA": Largest Algebraic (largest values)

    • "SA": Smallest Algebraic (smallest values)

  • v0 – Optional starting vector of shape (n,). None uses the deterministic all-ones start vector.

  • ncv (int | None) – Number of Lanczos basis vectors to build before extracting Ritz pairs. A larger value improves accuracy at the cost of more memory. Defaults to max(2*k+1, k+1).

  • maxiter (int | None) – Not yet supported because the current implementation performs one ncv-bounded Ritz extraction, not an implicitly restarted convergence loop. Pass None (the default).

  • tol (float) – Not yet supported for the same reason. Pass 0.0 (the default).

  • return_eigenvectors (bool) – When True (the default), return both eigenvalues and eigenvectors. When False, return only the eigenvalues.

Returns:

When return_eigenvectors=True, a tuple (values, vectors) where values has shape (k,) and vectors has shape (n, k). When return_eigenvectors=False, returns values alone.

Raises:
mlx_sparse.linalg.eigs(A, k=6, *, which='LM', v0=None, ncv=None, maxiter=None, tol=0.0, return_eigenvectors=True)[source]#

Compute a few eigenpairs of a general sparse square matrix.

Uses the native CSR Arnoldi-based eigensolver (an implicitly restarted Arnoldi method) to find the k eigenpairs of the general (possibly non-symmetric) sparse matrix A that match the criterion specified by which. Each Arnoldi step dispatches a sparse matrix-vector product to the GPU via the native Metal kernel.

For symmetric matrices, eigsh() is faster and more accurate because it uses the symmetric Lanczos recurrence instead of the full Arnoldi factorization.

GPU note:

When GPU execution is selected, Arnoldi factorisation uses the native Arnoldi kernel. The small Hessenberg eigensolve, Ritz value selection, and output vector assembly run on the CPU after the basis and Hessenberg matrix are copied back to host memory.

Parameters:

  • A – Sparse square matrix. Must be a CSRArray, COOArray, or CSCArray. Float16 and bfloat16 inputs are promoted to float32.

  • k (int) – Number of eigenpairs to compute. Must satisfy 0 < k < A.shape[0]. Defaults to 6.

  • which (str) –

    Which eigenpairs to return. Accepted values:

    • "LM": eigenvalues Largest in Magnitude (default)

    • "SM": eigenvalues Smallest in Magnitude

    • "LR": Largest Real part

    • "SR": Smallest Real part

  • v0 – Optional starting vector of shape (n,). None uses the deterministic all-ones start vector.

  • ncv (int | None) – Dimension of the Arnoldi factorization before restart. Defaults to max(2*k+1, k+1).

  • maxiter (int | None) – Not yet supported because the current implementation performs one ncv-bounded Ritz extraction, not an implicitly restarted convergence loop. Pass None (the default).

  • tol (float) – Not yet supported for the same reason. Pass 0.0 (the default).

  • return_eigenvectors (bool) – When True (the default), return both eigenvalues and eigenvectors. When False, return only the eigenvalues.

Returns:

When return_eigenvectors=True, a tuple (values, vectors) where values has shape (k,) and vectors has shape (n, k). When return_eigenvectors=False, returns values alone.

Raises:
mlx_sparse.linalg.svds(A, k=6, *, which='LM', v0=None, ncv=None, maxiter=None, tol=0.0, return_singular_vectors=True)[source]#

Compute a few singular triplets of a sparse matrix.

Uses the native CSR Lanczos iteration applied to the normal operator A.T @ A to find the k singular triplets (left singular vectors, singular values, and right singular vectors) of the sparse matrix A that match the criterion specified by which.

GPU note:

When GPU execution is selected, the normal-operator Lanczos recurrence uses a dedicated native A.T @ (A @ v) path. The two sparse products are kept inside one native step and the intermediate A @ v vector is not materialized on the host. The small tridiagonal eigensolve, Ritz vector back transformation, and returned singular-vector assembly still run on the CPU after the Lanczos basis is synchronized.

Parameters:

  • A – Sparse matrix of shape (m, n). Must be a CSRArray, COOArray, or CSCArray. Float16 and bfloat16 inputs are promoted to float32.

  • k (int) – Number of singular triplets to compute. Must satisfy 0 < k < min(A.shape). Defaults to 6.

  • which (str) –

    Which singular values to return. Accepted values:

    • "LM": Largest in Magnitude (default)

    • "SM": Smallest in Magnitude

  • v0 – Optional starting vector for the right singular-vector Krylov basis, with shape (A.shape[1],). None uses the deterministic all-ones vector.

  • ncv (int | None) – Number of Lanczos basis vectors to build. Defaults to max(2*k+1, k+1).

  • maxiter (int | None) – Not yet supported because the current implementation performs one ncv-bounded normal-operator Ritz extraction. Pass None (the default).

  • tol (float) – Not yet supported for the same reason. Pass 0.0 (the default).

  • return_singular_vectors (bool | str) –

    Controls which vectors are returned.

    • True: return (u, s, vh) (default)

    • False: return only s (the singular values)

    • "u": return (u, s, None)

    • "vh": return (None, s, vh)

Returns:

When return_singular_vectors=True, a tuple (u, s, vh) where u has shape (m, k), s has shape (k,), and vh has shape (k, n). See return_singular_vectors for other forms.

Raises:

  • NotImplementedError – If maxiter or tol are not at their default values.

  • ValueError – If k is out of range or return_singular_vectors is not a recognised value.

Sparse direct factorizations#

mlx_sparse.linalg.sparse_cholesky(A, *, upper=False)[source]#

Compute the sparse Cholesky factorization A = L @ L.T.

Performs a left-looking sparse Cholesky factorization on a real symmetric positive-definite (SPD) matrix stored in CSR, COO, or CSC format. COO/CSC inputs are converted once to canonical CSR before factorization. The resulting lower-triangular factor L is returned as a SparseCholesky object whose solve() method applies both triangular solves in sequence.

GPU note:

The factorization step runs on the CPU using the native sparse routine. The resulting SparseCholesky.solve dispatches triangular-solve kernels to the GPU when GPU execution is selected.

Parameters:

  • A – The matrix to factorize. Must be a CSRArray, COOArray, or CSCArray that is real, symmetric, and positive-definite. Float16 and bfloat16 inputs are promoted to float32 automatically.

  • upper (bool) – Not yet supported. Must be False (the default).

Returns:

A SparseCholesky dataclass holding the lower-triangular factor L as a CSRArray.

Raises:
Return type:

SparseCholesky

Example

Factorize a small SPD matrix and solve a linear system:

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

n = 8
L_sp = scipy.sparse.diags([-1, 4, -1], [-1, 0, 1],
                          shape=(n, n), format='csr').astype(np.float32)
A = ms.csr_array(
    (mx.array(L_sp.data), mx.array(L_sp.indices), mx.array(L_sp.indptr)),
    shape=L_sp.shape, canonical=True,
)
factor = linalg.sparse_cholesky(A)
b = mx.ones((n,), dtype=mx.float32)
x = factor.solve(b)
mlx_sparse.linalg.cholesky(A, *, upper=False)[source]#

Alias for sparse_cholesky().

Provided for compatibility with SciPy-style naming. All arguments and return values are identical to sparse_cholesky().

GPU note:

The factorization step runs on the CPU. The returned factor uses GPU triangular-solve kernels when GPU execution is selected.

Parameters:
Returns:

A SparseCholesky holding the lower-triangular factor L.

Return type:

SparseCholesky

mlx_sparse.linalg.sparse_lu(A)[source]#

Compute the sparse LU factorization P @ A = L @ U.

Performs a sparse LU factorization with partial pivoting on a general (possibly non-symmetric) real square matrix stored in CSR, COO, or CSC format. COO/CSC inputs are converted once to canonical CSR before factorization. The row permutation P, unit lower-triangular factor L, and upper-triangular factor U are returned as a SparseLU object whose solve() method applies the full solve sequence.

GPU note:

The factorization step runs on the CPU using the native sparse routine. The resulting SparseLU.solve dispatches permutation and triangular-solve kernels to the GPU when GPU execution is selected.

Parameters:

A – The matrix to factorize. Must be a CSRArray, COOArray, or CSCArray that is real and non-singular. Float16 and bfloat16 inputs are promoted to float32 automatically.

Returns:

A SparseLU dataclass with fields perm (row permutation as an mlx.core.array), L (unit lower-triangular CSRArray), and U (upper-triangular CSRArray).

Raises:

TypeError – If A is a dense array or has an unsupported dtype.

Return type:

SparseLU

Example

Factorize a non-symmetric sparse matrix and solve a system:

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

n = 8
rng = np.random.default_rng(0)
B = scipy.sparse.random(n, n, density=0.4, format='csr',
                        dtype=np.float32, random_state=rng)
B = B + scipy.sparse.eye(n, dtype=np.float32) * n
A = ms.csr_array(
    (mx.array(B.data), mx.array(B.indices), mx.array(B.indptr)),
    shape=B.shape, canonical=True,
)
factor = linalg.sparse_lu(A)
b = mx.ones((n,), dtype=mx.float32)
x = factor.solve(b)
mlx_sparse.linalg.splu(A)[source]#

Alias for sparse_lu().

Provided for compatibility with SciPy-style naming. All arguments and return values are identical to sparse_lu().

GPU note:

The factorization step runs on the CPU. The returned factor uses GPU permutation and triangular-solve kernels when GPU execution is selected.

Parameters:

A – Matrix to factorize in CSRArray, COOArray, or CSCArray format.

Returns:

A SparseLU dataclass with fields perm, L, and U.

Return type:

SparseLU

mlx_sparse.linalg.factorized(A, *, method='auto')[source]#

Factorize a sparse matrix once and return a reusable solve object.

method="auto" chooses LU for square matrices and QR for rectangular matrices. On Accelerate-enabled Apple builds, supported real float32 direct solves use opaque Accelerate factorization objects. Without Accelerate, square LU and Cholesky methods fall back to the existing native explicit sparse factors.

Parameters:

  • A – Coefficient matrix in CSR, CSC, or COO format. Float16 and bfloat16 inputs are promoted to float32 automatically.

  • method (str) – One of "auto", "lu", "cholesky", "ldlt", "qr", or "cholesky_ata". "ldlt", "qr", and "cholesky_ata" require an Accelerate-enabled Apple build.

Returns:

A FactorizedSolve object with solve() and __call__ methods.

Raises:

  • TypeError – If A is dense or has an unsupported dtype.

  • ValueError – If the selected method is incompatible with A.shape.

  • NotImplementedError – If the selected method requires Accelerate but this build does not provide it.

Return type:

FactorizedSolve

mlx_sparse.linalg.spsolve(A, b)[source]#

Solve the sparse linear system A @ x = b directly.

Computes a direct sparse factorization of A and immediately applies it to b. Accelerate-enabled Apple builds transparently use the Accelerate sparse LU path for supported real square systems. Other builds and unsupported cases use the native sparse_lu() path.

For repeated solves with the same A but different right-hand sides, call factorized() once and reuse the resulting FactorizedSolve object to avoid re-factorizing.

GPU note:

Direct factorization runs on the CPU. The native fallback row permutation and triangular solves use GPU kernels when GPU execution is selected. The Accelerate fast path uses Apple’s CPU sparse solver.

Parameters:

  • A – Coefficient matrix. Must be a CSRArray, COOArray, or CSCArray that is real and non-singular.

  • b – Right-hand side vector of shape (n,).

Returns:

Solution vector x of shape (n,) as an mlx.core.array.

Raises:

  • TypeError – If A is a dense array or has an unsupported dtype.

  • ValueError – If A is not square or b has an incompatible shape.

Return type:

mlx.core.array

class mlx_sparse.linalg.FactorizedSolve(_solver, shape, method, backend, rhs_size, solution_size)[source]#

Bases: object

Reusable sparse solve object.

FactorizedSolve intentionally exposes only solve behavior and metadata, not explicit sparse factors. Accelerate-enabled Apple builds use opaque Accelerate factorization objects for supported methods. Other supported square methods fall back to the existing native explicit-factor path.

Parameters:
shape: tuple[int, int]#
method: str#
backend: str#
rhs_size: int#
solution_size: int#
solve(b)[source]#

Solve for one or more right-hand sides using the stored factorization.

Parameters:

b – Right-hand side array of shape (rhs_size,) or (rhs_size, nrhs).

Returns:

Solution array of shape (solution_size,) or (solution_size, nrhs).

Return type:

mlx.core.array

__call__(b)[source]#

Alias for solve().

Return type:

mlx.core.array

class mlx_sparse.linalg.SparseCholesky(L)[source]#

Bases: object

Sparse Cholesky factorization A = L @ L.T.

The factor L is stored as a mlx_sparse.CSRArray. Numeric factorization is performed by the native sparse left-looking routine and the solves use sparse triangular CSR kernels.

GPU note:

The numeric factorization that creates L runs on the CPU. Calls to solve() use native triangular-solve kernels, which run on the GPU when GPU execution is selected.

Parameters:

L (CSRArray)

L: CSRArray#
property shape: tuple[int, int]#

Shape of the factored square matrix.

solve(b)[source]#

Solve A @ x = b using the stored Cholesky factor.

Performs two sparse triangular solves: a forward solve with L followed by a backward solve with L.T. Both steps use native CSR triangular-solve kernels.

GPU note:

When GPU execution is selected, both triangular solves dispatch native GPU kernels. The Python method call and shape checks run on the host.

Parameters:

b – Right-hand side array of shape (n,) or (n, nrhs).

Returns:

Solution array of shape (n,) or (n, nrhs) as an mlx.core.array.

Raises:

ValueError – If b has the wrong shape.

Return type:

mlx.core.array

__call__(b)[source]#

Alias for solve(). Allows the factorization to be called directly as factor(b).

Return type:

mlx.core.array

class mlx_sparse.linalg.SparseLU(perm, L, U)[source]#

Bases: object

Sparse LU factorization P @ A = L @ U.

L and U are CSR sparse factors. perm stores the row permutation applied before factorization.

GPU note:

The numeric LU factorization that creates perm, L, and U runs on the CPU. Calls to solve() use native permutation and triangular-solve kernels, which run on the GPU when GPU execution is selected.

Parameters:
perm: mlx.core.array#
L: CSRArray#
U: CSRArray#
property shape: tuple[int, int]#

Shape of the factored square matrix.

solve(b)[source]#

Solve A @ x = b using the stored LU factors.

Applies the row permutation P, then performs a forward solve with the unit lower-triangular factor L and a backward solve with the upper-triangular factor U. All steps use native CSR kernels.

GPU note:

When GPU execution is selected, the row permutation and both triangular solves dispatch native GPU kernels. The Python method call and shape checks run on the host.

Parameters:

b – Right-hand side array of shape (n,) or (n, nrhs).

Returns:

Solution array of shape (n,) or (n, nrhs) as an mlx.core.array.

Raises:

ValueError – If b has the wrong shape.

Return type:

mlx.core.array

__call__(b)[source]#

Alias for solve(). Allows the factorization to be called directly as factor(b).

Return type:

mlx.core.array

Operators and sparse reductions#

class mlx_sparse.linalg.LinearOperator(shape, matvec=None, *, matvec_fn=None, dtype=None, matmat=None, matmat_fn=None, rmatvec=None, rmatvec_fn=None, _sparse_array=None)[source]#

Bases: object

Sparse/matrix-free operator interface.

This class stores callables only. It does not densify an operator or provide dense fallbacks. Sparse arrays are normalized to canonical CSR and wrapped with native CSR matvec/matmat kernels through aslinearoperator().

Parameters:

  • matvec (Matvec | None)

  • matvec_fn (Matvec | None)

  • matmat (Matmat | None)

  • matmat_fn (Matmat | None)

  • rmatvec (Matvec | None)

  • rmatvec_fn (Matvec | None)

shape#
matvec_fn#
dtype#
matmat_fn#
rmatvec_fn#
property ndim: int#

Number of dimensions exposed by every linear operator.

matvec(x)[source]#

Apply the operator to a vector: compute A @ x.

Parameters:

x – Input vector of shape (n,).

Returns:

Output vector of shape (m,) as an mlx.core.array.

Raises:

ValueError – If x is not rank-1 or has the wrong length.

Return type:

mlx.core.array

matmat(X)[source]#

Apply the operator to a matrix: compute A @ X.

Parameters:

X – Input matrix of shape (n, k).

Returns:

Output matrix of shape (m, k) as an mlx.core.array.

Raises:

  • NotImplementedError – If no matmat_fn was provided at construction time.

  • ValueError – If X is not rank-2 or has the wrong leading dimension.

Return type:

mlx.core.array

rmatvec(x)[source]#

Apply the adjoint operator to a vector: compute A.H @ x.

Parameters:

x – Input vector of shape (m,).

Returns:

Output vector of shape (n,) as an mlx.core.array.

Raises:
Return type:

mlx.core.array

__matmul__(rhs)[source]#

Apply the operator to a vector or dense matrix with @.

property T: LinearOperator#

Transpose operator. (op.T) @ x computes A.T @ x.

For real operators A.T == A.H. For complex operators the formula A.T @ x = conj(A.H @ conj(x)) is used so no extra kernel is needed. Requires rmatvec_fn to be defined.

property H: LinearOperator#

Hermitian (conjugate) transpose operator. (op.H) @ x computes A.H @ x.

Requires rmatvec_fn to be defined (which stores A.H). The double adjoint (A.H).H recovers the original A.

mlx_sparse.linalg.aslinearoperator(A)[source]#

Wrap a sparse matrix or callable as a LinearOperator.

Accepts several input forms and returns a LinearOperator that exposes matvec(), matmat(), and rmatvec() via the native sparse kernels where possible.

Parameters:

A –

The object to wrap. Accepted types:

  • LinearOperator: returned unchanged.

  • CSRArray, COOArray, or CSCArray: converted once to canonical CSR and wrapped with native CSR matvec/matmat/rmatvec kernels.

  • SciPy sparse matrix (scipy.sparse): converted to CSR via from_scipy() then wrapped.

  • (shape, matvec) or (shape, matvec, matmat) tuple: the callables are stored directly with no conversion.

Returns:

A LinearOperator instance.

Raises:

TypeError – If A is not one of the accepted types.

Return type:

LinearOperator

mlx_sparse.linalg.dot(a, b)[source]#

Compute the Frobenius dot product of two sparse matrices.

Returns sum(a * b) over all stored non-zero pairs (no conjugation), using the native CSR sorted-merge kernel for efficient sparse-sparse element-wise accumulation.

Parameters:
Returns:

A scalar mlx.core.array equal to sum(a * b).

Raises:

TypeError – If a or b is not a supported sparse type.

mlx_sparse.linalg.vdot(a, b)[source]#

Compute the Frobenius inner product of two sparse matrices.

Returns sum(conj(a) * b) over all stored non-zero pairs, using the native CSR sorted-merge kernel for efficient sparse-sparse element-wise accumulation. Equivalent to dot(conj(a), b) for real matrices.

Parameters:
Returns:

A scalar mlx.core.array equal to sum(conj(a) * b).

Raises:

TypeError – If a or b is not a supported sparse type.