Skip to content

Linear solvers¤

ResistanceDistance and RSPDistance internally solve one or more sparse linear systems. The choice of solver directly affects both runtime and memory usage.

By default, when no solver is provided, JAXScape falls back to dense matrix inversion, which is only suitable for small graphs. For any non-trivial landscape, passing an explicit solver is strongly recommended. CholmodSolver is the preferred choice when memory permits.

lineax solvers¤

JAXScape accepts any Lineax solver. For example, to use the Conjugate Gradient method:

import lineax as lx
from jaxscape import ResistanceDistance

solver = lx.CG(rtol=1e-6, atol=1e-6)
distance = ResistanceDistance(solver=solver)

See the Lineax documentation for the full list of available solvers and their options.

JAXScape solvers¤

JAXScape provides three optional high-performance solvers as extras:

Solver GPU Memory cost Pros / Cons
CholmodSolver CPU only. Uses an external host callback, so accelerator-backed arrays must cross the host-device boundary. High Fastest option on CPU for a fixed sparse SPD system when memory permits. Handles multiple right-hand sides efficiently once inside the direct solve. Main downside is factorization memory, plus the callback introduces a data-transfer bottleneck and prevents a fully on-device GPU solve path.
PyAMGSolver CPU only. Built around SciPy CG + PyAMG through an external host callback. Moderate Lower memory than direct Cholesky and useful when a direct factorization is too expensive. Main downside is that each solve crosses into Python/NumPy space; batched solves are handled by a Python loop over right-hand sides, so throughput is usually weaker than the other two solvers.
AMJaxCGSolver Yes after initialization. The one-time hierarchy build uses PyAMG on CPU, but the initialized CG + AMJax preconditioner solve runs in JAX and can execute on CPU or GPU. Moderate Best choice for repeated, JIT-compiled, differentiable solves. Reuses initialized preconditioner state across vmapped batched solves without a per-solve host callback. The preconditioner can be reused while CG state is refreshed for related operators with the same structure.

Installation: 

uv add jaxscape --extra cholespy   # Cholesky solver
uv add jaxscape --extra pyamg      # PyAMG solver
uv add jaxscape --extra amjax      # Lineax CG + AMJax preconditioner

CI/CD coverage

These optional solvers are not included in the standard CI test suite.

Algebraic multigrid solver¤

AMJaxCGSolver separates hierarchy construction from the actual iterative solve. This matters when you want the solve itself to remain JIT-compatible and differentiable.

In algebraic multigrid, the preconditioner is built as a hierarchy of progressively coarser linear systems derived from the original sparse operator. This hierarchy consists of coarse operators together with prolongation and restriction maps between levels; one multigrid cycle smooths the error on the fine level, transfers the residual to coarser levels, approximately solves there, and interpolates the correction back to the fine level.

For direct sparse solves against an unchanged matrix, initialize the full solver state once:

import jax.numpy as jnp
from jax.experimental.sparse import BCOO
from jaxscape.solvers import AMJaxCGSolver, BCOOLinearOperator, linear_solve

A = BCOO.from_scipy_sparse(...)
b = jnp.ones(A.shape[0], dtype=A.data.dtype)

solver = AMJaxCGSolver(rtol=1e-6, atol=1e-6, max_steps=1_000)
state = solver.init(BCOOLinearOperator(A), {})
x = linear_solve(A, b, solver, state=state)

For related matrices whose values change but whose geometry is stable, initialize only the preconditioner and let JAXScape materialize fresh CG state for each current operator:

solver = AMJaxCGSolver(rtol=1e-6, atol=1e-6, max_steps=1_000)
preconditioner_state = solver.init_preconditioner(BCOOLinearOperator(A0), {})

x = linear_solve(A1, b, solver, state=preconditioner_state)

For resistance distance, initialize preconditioner state against a reference graph so the AMJax hierarchy is built on the grounded Laplacian used internally:

from jaxscape import GridGraph, ResistanceDistance
from jaxscape.solvers import AMJaxCGSolver

grid = GridGraph(permeability, fun=lambda x, y: (x + y) / 2)
distance = ResistanceDistance(
  solver=AMJaxCGSolver(rtol=1e-6, atol=1e-6, max_steps=1_000)
)
state = distance.init_preconditioner(grid)

R = distance(grid, state=state)

Use distance.init(grid) instead when the graph is fixed and you want to reuse the full solver state.

jaxscape.solvers.cholmodsolver.CholmodSolver ¤

A linear solver that uses CHOLMOD (via cholespy) to solve a sparse linear system. Uses direct Cholesky factorization for symmetric positive definite matrices.

Example

from jaxscape.solvers import CholmodSolver

solver = CholmodSolver()
distance = ResistanceDistance(solver=solver)
state = distance.init(grid)  # pre-factorizes the grounded Laplacian
dist = distance(grid, state=state)

Warning

cholespy must be installed to use this solver.

Warning

Float64 callback-backed solves require process-wide x64 support to be enabled before JAX work starts, e.g. via JAX_ENABLE_X64=1 or jax.config.update("jax_enable_x64", True). A thread-local with jax.enable_x64() block is not sufficient for host callback threads, and can surface as a callback dtype mismatch such as Expected: float64, Actual: float32.

jaxscape.solvers.pyamgsolver.PyAMGSolver ¤

A linear solver that uses PyAMG to solve a sparse linear system.

Example

from jaxscape.solvers import PyAMGSolver

solver = PyAMGSolver(rtol=1e-6, atol=1e-6, maxiter=100_000)
distance = ResistanceDistance(solver=solver)
dist = distance(grid)

# Custom AMG method
import pyamg
solver = PyAMGSolver(pyamg_method=pyamg.ruge_stuben_solver)

Warning

PyAMG must be installed to use this solver.

jaxscape.solvers.amjaxcgsolver.AMJaxCGSolver ¤

A Lineax CG solver with an AMJax algebraic multigrid preconditioner.

The PyAMG hierarchy setup is performed once before entering Lineax's traced solve path. The resulting AMJax V-cycle is wrapped as a lineax.FunctionLinearOperator and supplied to lineax.CG through the preconditioner option.

Example

from jaxscape.solvers import AMJaxCGSolver, BCOOLinearOperator, linear_solve

solver = AMJaxCGSolver(rtol=1e-6, atol=1e-6, max_steps=1_000)
state = solver.init(BCOOLinearOperator(A), {})
x = linear_solve(A, b, solver, state=state)

Warning

amjax and pyamg must be installed to use this solver.

Reverse-mode gradients through solver calls require a symmetric AMJax preconditioner. This holds for the default matching pre/post Jacobi smoothers, but may not hold for custom non-symmetric smoother choices. Such configurations can still be valid for forward solves, but Lineax may need to transpose the preconditioner during autodiff.

Advanced: BCOOLinearOperator¤

JAXScape exposes a lineax-compatible linear operator that wraps JAX's native BCOO sparse matrix format. This allows any Lineax solver to operate directly on sparse matrices without converting to a dense representation, and is used internally by CholmodSolver, PyAMGSolver, and AMJaxCGSolver.

jaxscape.solvers.operator.BCOOLinearOperator ¤

lineax.MatrixLinearOperator wrapper for jax.experimental.sparse.BCOO matrices.