API Reference

Matrix or callable operator A(x). All matrices/operators are converted to jax.experimental.sparse.bcsr sparse matrices internally. In MPI mode this is the local partition.

Right-hand-side vector. In MPI mode this is the local RHS partition.

Optional initial guess (same shape as b; local partition in MPI mode). Defaults to zero. A good warm start (e.g. the previous solution in a time-stepping or optimization loop) cuts iterations; it does not change the converged solution or its gradients. Note that with the default RELATIVE_INI convergence the tolerance is relative to the initial residual, so a very good x0 tightens the target; consider convergence="ABSOLUTE" for warm-started loops.

AmgX configuration dictionary (see Solver Configuration for details). If None, JAX-AMG defaults are used.

Treat the matrix as a block matrix with square block_dim x block_dim blocks (e.g. coupled multi-component PDE systems with node-major interleaved unknowns: row i*block_dim + 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. Rows must be divisible by block_dim (each rank's local partition in MPI mode). Since AmgX's classical AMG does not support blocks, the AMG defaults switch to aggregation (SIZE_2 + BLOCK_JACOBI); explicitly configured CLASSICAL AMG is rejected. In MPI mode the aggregation defaults use block-Jacobi coarse sweeps instead of DENSE_LU_SOLVER (which is broken in AmgX for distributed block matrices on 3+ ranks and rejected if configured explicitly); see Solver Configuration.

MPI communicator (typically mpi4py.MPI.COMM_WORLD). If provided, the solve runs in MPI mode. If not provided, MPI mode can still be used if MPI metadata has already been attached via with_cache(..., mpi=...).

Global matrix row count for MPI mode. Required when comm is provided and MPI metadata is not pre-attached to A.

(row_start, row_end) owned by this rank in MPI mode. Required when comm is provided and MPI metadata is not pre-attached to A.

Optional file path to save detailed AmgX solver statistics. If None, no file is created.

For repeated solves with the same sparsity pattern, skip warm AMGX_solver_resetup and keep the cached hierarchy. This is cheaper per solve but may require more iterations if matrix coefficients change significantly.

Basis of null(A) for singular systems: "constant", a length-n vector, or an (n, k) array (local rows in MPI mode). The solution is pinned orthogonal to it, and the transpose of that pin projects the adjoint right-hand side onto range(Aᵀ), which makes the backward solve converge. Gradients w.r.t. A assume perturbations that preserve the declared null spaces (dA·N = 0, Mᵀ·dA = 0), as coefficient changes of a conservative discretization do. Defaults to the basis attached with with_cache.

Basis of null(Aᵀ) (same formats). b is projected onto range(A) (removed fraction in info["rhs_inconsistency"]) and, by transposition, the adjoint solution is pinned: the forward returns A⁺b, jax.grad returns (Aᵀ)⁺g. Equals nullspace for symmetric A (filled in when A is marked symmetric); for nonsymmetric A it differs, e.g. A = D⁻¹L has nullspace="constant" but transpose_nullspace=V (cell volumes). Defaults to the basis attached with with_cache.

Additional AmgX config parameters. These override values in config when both are provided.

Solution vector (float32 or float64). In MPI mode, returns local portion.

Dictionary containing iterations, residual, status, and residual_history (residual norm per outer iteration, entry 0 being the initial residual; inside jit it has fixed length max_iters + 1 with NaN padding past entry iterations). With transpose_nullspace, also rhs_inconsistency (‖b − b'‖/‖b‖).

A·1 = 0 without a declared nullspace; nullspace without transpose_nullspace (or vice versa) for a matrix not marked symmetric; a basis failing A·N ≈ 0 / Aᵀ·M ≈ 0 (the latter not checked in MPI mode); or a DENSE_LU_SOLVER coarse solve. Checks run only on concrete matrix values.

This process's CSR row partition with shape (n_local, n_global) and global column indices. A matrix-free operator is also accepted, provided it carries cached coloring information (jaxamg.with_cache(op, coloring=...)) so its shape and sparsity pattern are known; it is materialized once here. Null-space bases attached with with_cache(A_local, nullspace=..., transpose_nullspace=...) (local rows) are kept and applied by every solve.

Global row-sharded RHS used to validate the matrix partition, mesh, and numerical dtype.

MPI communicator spanning every JAX process, with rank order matching the JAX process order. Defaults to MPI.COMM_WORLD.

One-dimensional JAX device mesh. If omitted, use the mesh from b.sharding.

Name of the mesh axis that partitions rows and packed values.

class:ShardedMatrix whose data attribute contains the global

sharded values. The original global matrix is never materialized.

This rank's unpadded values with shape (n_local,).

MPI communicator spanning every JAX process, with rank order matching mesh. Defaults to MPI.COMM_WORLD.

One-dimensional JAX device mesh with one device per MPI rank. Defaults to a mesh over the first comm.size JAX devices (all devices in a typical multi-process job).

Optional true global length. When provided, it is checked against the sum of local lengths.

Mesh axis used to partition the vector.

A global JAX array whose axis uses P(axis_name). Its physical

length is comm.size * max(local_sizes); values are cast to

float32 unless already float32 or float64.

Distributed matrix created with :func:make_sharded_matrix.

Global row-sharded JAX vector. Construct it with :func:make_sharded_vector, which handles unequal row counts.

AmgX configuration. Defaults to the JAX-AMG MPI configuration.

Whether the global matrix is symmetric. When False (the default), the distributed transpose is prepared once during solver creation for the reverse-mode adjoint.

AmgX block dimension. The matrix retains its scalar CSR representation, and every local row partition must be divisible by this value.

Reuse the cached AmgX hierarchy across solves with the same sparsity pattern.

Prepare the AmgX configuration with solver-statistics output enabled so a later direct call with solver(..., save_stats_file=...) produces a complete stats file.

A callable solver(b, x0=None, *, A=None, save_stats_file=None).

A is this rank's operator or matrix with the sparsity fixed here

(its closed-over parameters must be identical on every rank) or the

packed global values A.data; it may be omitted for a direct call

but is required under a JAX transform. A.local_matrix(gradient)

unpacks a packed matrix gradient and solver.local_vector(value)

removes vector padding. Info values have one entry per rank.

A matrix or operator.

Cached coloring information from cache_coloring().

Cached MPI metadata from cache_mpi_metadata().

If True, indicates the matrix is symmetric, allowing optimizations like skipping transpose in backward pass.

Default for jaxamg.solve's nullspace ("constant", a vector, or an (n, k) array; local rows in MPI mode).

Default for jaxamg.solve's transpose_nullspace.

The same matrix/operator with requested cache attached.

A callable operator A(x) that returns A @ x.

Shape of the operator (n, m) or int size (for an n×n matrix). For a distributed operator this is the local block (n_local, n_global).

Cached coloring information for reattachment with with_cache(..., coloring=...).

AmgX configuration dict or string

MPI communicator (from mpi4py.MPI.COMM_WORLD)

Global matrix size (total rows across all ranks)

tuple (row_start, row_end) indicating which rows this rank owns

Matrix or operator to compute max nnz for buffer sizing

If True, the backward pass never transposes, so the transpose output size (nnz_out) is left unset (None). Should match the is_symmetric passed to with_cache; the default (False) computes it, which is always safe.

If True, prepare the config with solver statistics output enabled, so a later solve(..., save_stats_file=...) on the cached matrix produces a complete stats file.

BSR block size for AmgX (see jaxamg.solve). Each rank's local partition must be divisible by it.

Use the singular-system AMG defaults (see Solver Configuration). Implied when null-space bases are already attached to A via with_cache.

A dictionary containing MPI metadata.

Matrix or callable operator to precondition.

Optional AmgX configuration. If omitted, a single-cycle AMG approximate-inverse config is used.

Optional MPI communicator for distributed solves. If A already has MPI metadata attached via jaxamg.with_cache(..., mpi=...), this may be omitted.

Global matrix row count for MPI mode.

Local row partition (row_start, row_end) for MPI mode.

Optional stats output path passed to jaxamg.solve(...).

If True, the returned callable yields (x, info) instead of only x.

Additional solver config overrides.

A callable representing an approximate inverse M^{-1}.

The system operator to precondition.

Optional AmgX configuration (see make_preconditioner).

Lineax tags for the returned preconditioner. By default the operator's own tags are inherited (A⁻¹ shares A's symmetry/definiteness), which CG needs in order to accept the preconditioner. Pass an explicit value (e.g. ()) to override.

Optional MPI communicator for distributed solves.

Global matrix row count for MPI mode.

Local row partition (row_start, row_end) for MPI mode.

Optional stats output path passed to jaxamg.solve(...).

Additional solver config overrides forwarded to make_preconditioner.

A lineax.FunctionLinearOperator approximating A⁻¹.

A dictionary with cache size/capacity and entry summaries

for single-GPU and MPI caches, plus whether isolated mode

(JAXAMG_CACHE_SIZE=0) is active.