Add optional JAX backend for GPU-accelerated CG and MINRES solvers#19

Open

tschm with Copilot wants to merge 9 commits into

mainfrom

copilot/add-jax-backend-support

Commits on May 1, 2026

Initial plan
Copilot
authored
Add JAX backend support to Problem class with solve_cg dispatch

Copilot
and
tschm
authored
Plan: add JAX MINRES port for solve_minres JAX backend

Copilot
and
tschm
authored
Add JAX MINRES port: _minres_jax function and _solve_minres_jax method

Copilot
and
tschm
authored
Fix CG iter count with _cg_jax helper; add benchmarks/jax_backend.py

Copilot
and
tschm
authored
Refactor: _jax_arrays helper, parametrize validation test, clean up imports

Copilot
and
tschm
authored
Fix na_str parameter name in benchmark _fmt helper

Copilot
and
tschm
authored
Fix JAX Krylov performance: lax.while_loop MINRES + jax.scipy.linalg.cg, clearer benchmark CPU warning

Copilot
and
tschm
authored
Add jax-metal to jax extra; extend benchmark to N=5000

tschm
and
claude
committed