Sparse linalg: spectral routines#

mlx_sparse.linalg provides three native sparse spectral routines that compute only the k extreme eigenpairs or singular triplets, not the full decomposition. This is the main advantage over dense mx.linalg.eigh / mx.linalg.svd, which always compute all values.

Routine

Method

Matrix

Returns

eigsh

Native CSR Lanczos

Hermitian/symmetric

(vals, vecs)

eigs

Native CSR Arnoldi

Any square

(vals, vecs)

svds

Normal-operator Lanczos

Any rectangular

(u, s, vh)

The which parameter selects which extremal eigenvalues to compute: "LM" (Largest Magnitude, default) or "SM" (Smallest Magnitude).

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

Test matrix: sparse graph Laplacian#

The graph Laplacian of a path graph on n nodes is a tridiagonal symmetric positive semi-definite matrix. Its eigenvalues are 2 - 2*cos(π*k/n) for k = 0, …, n-1, giving a known reference.

import scipy.sparse

n = 20
# Path-graph Laplacian: tridiagonal with 2 on diagonal and -1 on off-diagonals
lap = scipy.sparse.diags(
    [[-1.0]*(n-1), [2.0]*n, [-1.0]*(n-1)],
    [-1, 0, 1], shape=(n, n), dtype=np.float32,
).tocsr()
# The smallest eigenvalue is 0 (constant eigenvector), add small shift to make it PD
lap_pd = lap + scipy.sparse.eye(n, dtype=np.float32) * 0.1

A = ms.csr_array(
    (mx.array(lap_pd.data), mx.array(lap_pd.indices), mx.array(lap_pd.indptr)),
    shape=(n, n), canonical=True,
)
print(f"shape={A.shape}, nnz={A.nnz}")

# Dense reference eigenvalues
ref_vals = np.linalg.eigvalsh(lap_pd.toarray())
print(f"True eigenvalues (sorted): {np.round(np.sort(ref_vals), 4)}")

shape=(20, 20), nnz=58
True eigenvalues (sorted): [0.1223 0.1889 0.2981 0.4475 0.6339 0.853  1.1    1.3693 1.655  1.9505
 2.2495 2.545  2.8307 3.1    3.347  3.5661 3.7525 3.9019 4.0111 4.0777]

eigsh: largest eigenvalues (Lanczos)#

eigsh(A, k=6, which='LM') runs the native Lanczos algorithm and returns the k eigenvalues of largest magnitude together with their eigenvectors. The ncv parameter (Krylov subspace size) controls accuracy vs. cost.

k = 4
vals_lm, vecs_lm = linalg.eigsh(A, k=k, which="LM", ncv=30)
mx.eval(vals_lm, vecs_lm)

print(f"Largest {k} eigenvalues (Lanczos): {np.sort(np.array(vals_lm))[::-1]}")
print(f"True largest {k}: {np.sort(ref_vals)[::-1][:k]}")

Largest 4 eigenvalues (Lanczos): [4.0776653 4.011147  3.9019392 3.7524784]
True largest 4: [4.0776615 4.0111456 3.9019377 3.7524774]

# Verify eigenpairs: A @ v approx. λ * v for each returned pair
vals_np = np.array(vals_lm)
vecs_np = np.array(vecs_lm)

A_dense = lap_pd.toarray()
for i in range(k):
    v = vecs_np[:, i]
    lam = vals_np[i]
    residual = np.linalg.norm(A_dense @ v - lam * v)
    print(f"  λ={lam:.4f}  ||Av - λv||={residual:.2e}")

  λ=4.0777  ||Av - λv||=3.85e-06
  λ=4.0111  ||Av - λv||=1.72e-06
  λ=3.9019  ||Av - λv||=1.98e-06
  λ=3.7525  ||Av - λv||=1.45e-06

eigsh: smallest eigenvalues#

Pass which='SM' to request the smallest-magnitude eigenvalues. These correspond to the lowest-frequency modes of the graph Laplacian.

vals_sm, _ = linalg.eigsh(A, k=k, which="SM", ncv=30)
mx.eval(vals_sm)

print(f"Smallest {k} eigenvalues (Lanczos): {np.sort(np.array(vals_sm))}")
print(f"True smallest {k}: {np.sort(ref_vals)[:k]}")

Smallest 4 eigenvalues (Lanczos): [0.1223383  0.1888545  0.2980625  0.44752267]
True smallest 4: [0.12233825 0.18885429 0.29806218 0.44752234]

eigs: general sparse Arnoldi#

eigs uses the Arnoldi factorization and works for non-symmetric matrices. It returns complex eigenvalues in general; for symmetric matrices the imaginary parts are negligible.

vals_arnoldi, vecs_arnoldi = linalg.eigs(A, k=k, which="LM", ncv=30)
mx.eval(vals_arnoldi)

vals_real = np.array(vals_arnoldi)
print(f"Arnoldi eigenvalues: {np.round(np.sort(np.abs(vals_real))[::-1], 4)}")
print(f"Lanczos eigenvalues: {np.round(np.sort(np.abs(np.array(vals_lm)))[::-1], 4)}")

Arnoldi eigenvalues: [4.0775 4.0113 3.9019 3.7525]
Lanczos eigenvalues: [4.0777 4.0111 3.9019 3.7525]

svds: sparse singular value decomposition#

svds applies Lanczos to the normal operator A.T @ A without materializing the product. It accepts rectangular matrices and returns (u, s, vh) where u (left singular vectors), s (singular values), and vh (right singular vectors, transposed) satisfy u @ diag(s) @ vh ≈ A.

# Rectangular sparse matrix for svds
import scipy.sparse
m, nc = 30, 20
rng = np.random.default_rng(99)
rect = scipy.sparse.random(m, nc, density=0.15, format="csr", dtype=np.float32, random_state=rng)

B = ms.csr_array(
    (mx.array(rect.data), mx.array(rect.indices), mx.array(rect.indptr)),
    shape=(m, nc), canonical=True,
)
print(f"B shape={B.shape}, nnz={B.nnz}")

k_sv = 3
u, s, vh = linalg.svds(B, k=k_sv, which="LM", ncv=30)
mx.eval(u, s, vh)

s_ref = np.linalg.svd(rect.toarray(), compute_uv=False)
print(f"Sparse svds top-{k_sv}: {np.round(np.sort(np.array(s))[::-1], 4)}")
print(f"Dense SVD  top-{k_sv}: {np.round(s_ref[:k_sv], 4)}")

B shape=(30, 20), nnz=90
Sparse svds top-3: [2.6716 2.0926 1.9927]
Dense SVD  top-3: [2.6716 2.0926 1.9927]

# Verify singular triplets: B @ vh[i] approx. s[i] * u[i]
s_np = np.array(s)
u_np = np.array(u)
vh_np = np.array(vh)
B_dense = rect.toarray()

for i in range(k_sv):
    err = np.linalg.norm(B_dense @ vh_np[i] - s_np[i] * u_np[:, i])
    print(f"  s={s_np[i]:.4f}  ||B @ v_i - s_i * u_i||={err:.2e}")

  s=2.6716  ||B @ v_i - s_i * u_i||=1.87e-07
  s=2.0926  ||B @ v_i - s_i * u_i||=8.71e-08
  s=1.9927  ||B @ v_i - s_i * u_i||=5.36e-08

Comparing with dense: speedup depends on k/n#

Sparse spectral routines are advantageous when k << n. For large n and small k, the Lanczos/Arnoldi approach avoids the O(n³) cost of full dense decomposition. The crossover point depends on matrix sparsity and hardware; for k=4, n=20 the system is too small to see a speedup.

import time

# use CPU since MLX's eigh is not supported on GPU yet
ms.use_cpu()

def timer(fn, iters=10):
    for _ in range(3):
        mx.eval(fn())
    t0 = time.perf_counter()
    for _ in range(iters):
        mx.eval(fn())
    return 1000 * (time.perf_counter() - t0) / iters

A_dense_mx = mx.array(lap_pd.toarray())
sparse_ms = timer(lambda: linalg.eigsh(A, k=k, which="LM")[0])
dense_ms  = timer(lambda: mx.linalg.eigh(A_dense_mx)[0])

print(f"eigsh (k={k}, n={n}):  {sparse_ms:.3f} ms")
print(f"dense eigh (all n):  {dense_ms:.3f} ms")
print(f"speedup: {dense_ms/sparse_ms:.1f}x")

eigsh (k=4, n=20):  0.016 ms
dense eigh (all n):  0.086 ms
speedup: 5.3x