Solver Configuration

This page documents how to configure jaxamg.solve(...), including solver/preconditioner choices and common AmgX options.

Configuration input modes

jaxamg.solve(...) accepts configuration in three ways:

1. config=<dict>: a full configuration dictionary.
2. keyword arguments (solver=..., max_iters=..., etc.).
3. both together, where keyword arguments override values in config.

Example:

x, info = jaxamg.solve(
    A,
    b,
    config={"solver": "CG", "max_iters": 100},
    max_iters=50,  # overrides config["max_iters"]
)

Default config

If you do not provide a config, JAX-AMG uses:

{
    "solver": "PBICGSTAB",
    "preconditioner": {
        "solver": "AMG",
        "algorithm": "CLASSICAL",
        "selector": "PMIS",
        "interpolator": "D2",
        "smoother": {
            "solver": "BLOCK_JACOBI",
            "relaxation_factor": 0.9,
        },
        "presweeps": 1,
        "postsweeps": 1,
        "max_levels": 100,
        "strength_threshold": 0.5,
        "dense_lu_num_rows": 1,
        "aggressive_levels": 0,
        "coarse_solver": "DENSE_LU_SOLVER",
        "max_iters": 1,
        "cycle": "V",
    },
    "convergence": "RELATIVE_INI",
    "tolerance": 1e-6,
    "max_iters": 1000,
    "norm": "L2",
    "exact_coarse_solve": 1,
}

User-supplied keys are merged into these defaults, so you only need to specify what you want to change. For example, config={"preconditioner": {"solver": "AMG"}} inherits all the Classical AMG settings above.

JAX-AMG also enables residual tracking internally (monitor_residual=1, store_res_history=1) so info["residual"] and info["residual_history"] are always available. The history is the outer solver's convergence curve: entry i is the residual norm after outer iteration i, with entry 0 the initial residual. Outside jit it is trimmed to iterations + 1 entries; inside jit it has fixed length max_iters + 1, NaN-padded past entry iterations.

Block matrices

For coupled multi-component systems, pass block_dim=k to jaxamg.solve(...) to treat the matrix as having square k x k blocks (unknowns interleaved node-major: row i*k + c is component c of node i). A and b keep their ordinary scalar CSR/vector form — the conversion to AmgX's BSR format happens internally — and autodiff works unchanged. Rows must be divisible by block_dim (each rank's local partition in MPI mode).

Because AmgX's classical AMG does not support block matrices, the AMG defaults switch to aggregation AMG when block_dim > 1:

{
    "solver": "AMG",
    "algorithm": "AGGREGATION",
    "selector": "SIZE_2",
    "smoother": {"solver": "BLOCK_JACOBI", "relaxation_factor": 0.9},
    "presweeps": 1,
    "postsweeps": 1,
    "max_levels": 100,
    "min_coarse_rows": 32,
    "dense_lu_num_rows": 64,
    "coarse_solver": "DENSE_LU_SOLVER",
    "max_iters": 1,
    "cycle": "V",
}

In MPI mode the block defaults differ in one place: coarse_solver becomes {"solver": "BLOCK_JACOBI", "max_iters": 50} (and dense_lu_num_rows is dropped). AmgX's distributed DENSE_LU_SOLVER produces wrong coarse corrections for block matrices spanning 3 or more ranks — V-cycles silently diverge — so explicitly configuring DENSE_LU_SOLVER (as coarse solver or outer solver) together with comm and block_dim > 1 raises a ValueError.

Explicitly configuring CLASSICAL AMG together with block_dim > 1 raises a ValueError. Block-aware preconditioning is often the entire benefit: on a strongly coupled system, BLOCK_JACOBI under block_dim=2 inverts true 2x2 diagonal blocks (instead of scalar diagonal entries) and can converge in a handful of iterations where the scalar-preconditioned solve stalls.

MPI config

The communicator key can be set to MPI for standard CPU-based MPI or MPI_DIRECT for GPU-aware MPI. The default is MPI. To use GPU-aware MPI, ensure that your MPI installation supports it. For more details, see the MPI Guide.

config = {
    "solver": "PBICGSTAB",
    "preconditioner": {"solver": "MULTICOLOR_DILU"},
    "communicator": "MPI_DIRECT",
    "max_iters": 200,
    "tolerance": 1e-8,
}

Flat vs nested config

JAX-AMG supports both flat (simple) and nested config dict:

• Flat config (simple dict with top-level keys like solver, tolerance).

  config = {
      "solver": "CG",
      "preconditioner": {"solver": "JACOBI_L1"},
      "tolerance": 1e-8,
      "max_iters": 200,
  }
  

• Nested config (config_version: 2 with solver: {...} scope).

  config = {
      "config_version": 2,
      "solver": {
          "solver": "PCG",
          "preconditioner": {
              "solver": "AMG",
              "smoother": "JACOBI_L1",
          },
          "tolerance": 1e-6,
          "max_iters": 200,
      },
  }
  

Common configuration keys

These are commonly used and work in JAX-AMG:

Convergence criteria

AMG-specific keys

These go inside the AMG solver/preconditioner scope:

Supported Solvers and Preconditioners

AmgX supports a broader set of solvers, preconditioners, and smoothers: