| Jacobi | Yes | — | — |
| Gauss-Seidel | No | No | - Update of
x_i depends on x_j^{k+1} for j < i: irreducible sequential dependency lax.scan would produce an O(n) XLA graph, not vectorizable on GPU
|
| SOR | No | No | - Gauss-Seidel with omega: same sequential dependency
- Same
lax.scan limitation
|
| Block Gauss-Seidel | No | No | - Same structure as GS by blocks, sequential inter-block dependency
|
| Block Jacobi | No | Yes | - Blocks are independent: all can be solved simultaneously
- Block inverses precomputed at setup
|
| Schwarz (multiplicative) | No | No | - Each local solve modifies
x, creating a sequential intra-sweep dependency
|
| Polynomial (generic) | No | Yes | - General polynomial smoother:
p(A)r with static degree - Richardson and Chebyshev are special cases (see below)
|
| Richardson (via polynomial) | No | Yes | - One matrix-vector product and vector additions per iteration: fully parallel
omega precomputed at setup
|
| Chebyshev (via polynomial) | No | Yes | - Horner's rule with static degree: unrolled at tracing, each matmul parallel
- Coefficients precomputed from spectrum
|
| Jacobi-NE | No | Yes | - Same structure as Jacobi, fully parallel
D_{AA^H} precomputed at setup
|
| Gauss-Seidel-NE | No | No | - GS on normal equations
AA^Hy = b: sequential row by row
|
| Gauss-Seidel-NR | No | No | - GS on
A^HAx = A^Hb: sequential column by column
|
| Gauss-Seidel indexed | No | Partial | - Implementable if the index sets are independent (e.g. graph coloring)
- Otherwise same sequential dependency as GS
|
| CF-Jacobi / FC-Jacobi | No | Yes | - Jacobi on C/F-points separately
- Within each group, fully independent via
jnp.where on static mask
|
| CG as smoother | No | Yes | jax.scipy.sparse.linalg.cg with static maxiter is jit-compilable
|
| GMRES as smoother | No | Yes | jax.scipy.sparse.linalg.gmres with static maxiter is jit-compilable
|
| CGNE / CGNR as smoother | No | Yes (effort) | - Absent from
jax.scipy - Requires custom implementation (algorithmically simple: CG on normal equations)
|
| Jacobi | Yes | — | — |
| pinv | Yes | — | — |
| LU | Yes | — | — |
| QR | Yes | — | — |
| Cholesky | No | Yes | jax.scipy.linalg.cho_factor/cho_solve exist- Same pattern as LU already implemented
- Requires SPD matrix
|
| pinv2 | No | Yes | - pinv with
rcond parameter to control SVD truncation jnp.linalg.pinv already supports rcond
|
| splu (sparse LU) | No | No | - Relies on SuperLU (external C library): no native sparse LU in XLA
jax.pure_callback possible but breaks vmap and grad
|
| CG | No | Yes | jax.scipy.sparse.linalg.cg with static maxiter
|
| GMRES | No | Yes | jax.scipy.sparse.linalg.gmres with static max_iter
|
| BiCGSTAB | No | Yes | jax.scipy.sparse.linalg.bicgstab with static maxiter
|
| FGMRES | No | Yes (very high effort) | - Requires variable preconditioner per iteration
- Not supported by
jax.scipy
|
| V-cycle | Yes | — | — |
| W-cycle | Yes | — | — |
| F-cycle | Yes | — | — |
| AMLI | No | Partial | - Cycle itself is implementable
- Requires FGMRES as accelerator, not available in JAX
|
| CG | Yes | — | jax.scipy.sparse.linalg.cg(A, b, M=ml.aspreconditioner())
|
| GMRES | Yes | — | jax.scipy.sparse.linalg.gmres(A, b, M=ml.aspreconditioner())
|
| BiCGSTAB | Yes | — | jax.scipy.sparse.linalg.bicgstab(A, b, M=ml.aspreconditioner())
|
| FGMRES | No | Yes (very high effort) | - Absent from
jax.scipy - Requires variable preconditioner per Krylov iteration
|
| Ruge-Stuben | Yes | — | — |
| Smoothed Aggregation | Yes | — | — |
| Root Node | Yes | — | — |
| Adaptive SA | No | Yes (effort) | - Adaptive process to find near-null components
- Implementable with a fixed number of iterations unrolled at trace time
|
| AIR (Approx. Ideal Restriction) | No | Yes (effort) | - Requires solving local least-squares problems at setup
- Solvable with
jnp.linalg.lstsq if subdomain sizes are static
|
| Complex matrices | No | Yes | - JAX supports
complex64/complex128 natively - Requires adjustments to the current implementation
|
| Residual tracking under jit | No | No | lax.while_loop forbids side effects: list.append impossible inside XLA graph- Debug mode only via
jax.debug.callback, incompatible with vmap and grad
|