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.

What repeated solves with the same sparsity pattern and configuration keep from the cached solver. None (default) refreshes the matrix values and rebuilds the hierarchy (AMGX_solver_resetup). "hierarchy" refreshes the values but keeps the hierarchy: cheaper per solve, but it may need more iterations if the coefficients change significantly. "matrix" keeps both, for an operator whose values never change: A must be a callable declared with with_cache(A, pattern=...), and solves of operators with the same pattern object use the first solve's matrix. Inside jit, A is evaluated only when that matrix is not cached. Each kept matrix takes its own cache entry (JAXAMG_CACHE_SIZE); such solves run on a single GPU or through make_sharded_solver and are not differentiable.

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.

The pair (count, labels), as scipy.sparse.csgraph.connected_components returns it: a static label count and one integer label per row (local rows in MPI mode, possibly traced), -1 for rows in no component, for a disconnected domain: every declared nullspace/transpose_nullspace column then applies separately on each label's rows (the null space is block diagonal over the labels), e.g. nullspace="constant" with labels declares one constant vector per component. A label whose columns vanish on its rows (no rows, or all of them -1 in this solve) is skipped exactly; columns dependent within a label are rejected (concrete values). Without labels the null space is the columns themselves. Defaults to the labels attached with with_cache.

In MPI mode with labels, the reduction of each rank's per-label partial sums (an array of count rows) to the labels' totals, a linear function (a jaxamg.linear_map.LinearMap where JAX's own transpose would be wrong, as for an MPI all-reduce). By default, the all-reduce over global label numbers (one count-row array on every rank); a caller whose labels span ranks in its own numbering supplies the reduction. Defaults to the one attached with with_cache.

"implicit" (default) differentiates x = A⁻¹b implicitly on the fixed pattern, ẋ = A⁻¹(ḃ − Ȧx) with each A⁻¹ a native solve: forward and reverse mode, and higher orders where the operator supports them. "adjoint" is the reverse-only rule. Both use the same first reverse rule. Neither differentiates the AmgX iterations: derivatives are those of the exact solution, accurate to the solve tolerance.

Whether the solve declares a side effect to JAX (default True): every solve updates the solver cache (its setup for the sparsity pattern and configuration, possibly displacing another), which reuse="hierarchy" solves read. JAX keeps every such solve, and under jax.checkpoint (JAX 0.11.1 or later; older JAX rejects the effect) runs it once and keeps its output for the backward pass. False declares the solve pure: JAX may drop it when its result is unused or merge identical solves, and recomputes it under jax.checkpoint instead of keeping its output, which bounds the memory of checkpointed computations. Declare it only when no reuse="hierarchy" solve depends on this solve's setup, or when such a solve may use whichever setup is cached (its result then varies within the solve tolerance). Per-rank MPI solves keep mpi4jax's ordered effect either way, which jax.checkpoint rejects. See Caching.

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; or a DENSE_LU_SOLVER coarse solve. These checks run only on concrete matrix values.

The operator, linear in its inputs; it must not communicate.

The number of rows this rank owns.

Ascending global columns the rows reference outside this rank; x_ghost[k] holds column ghost_ids[k].

Optional declared jaxamg.Pattern over the combined (n_local, n_local + len(ghost_ids)) space. Without it the pattern is traced (else probed) on first use, which needs concrete parameters: derive per-step operators with HaloOperator.with_fn to reuse it inside jit and grad.

A HaloOperator.

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 and labels attached with with_cache(A_local, nullspace=..., transpose_nullspace=..., labels=..., label_sum=...) (local rows) are kept as every solve's defaults.

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.

JAX device mesh (see :func:make_sharded_vector). If omitted, use the mesh from b.sharding.

Mesh axis, or tuple of mesh axes in mesh order, that partitions rows and packed values (as for b).

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.

JAX device mesh with one device per MPI rank, in rank order. Its axis_name axes partition the rows; any other axis must have size one, so a program's multi-axis mesh (e.g. a Cartesian domain decomposition's) can be shared. Defaults to a one-dimensional 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, or tuple of mesh axes in mesh order, partitioning 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.

As in :func:jaxamg.solve. With "matrix", each solver keeps the matrix of its first solve, and inside jax.jit evaluates A only when that matrix is not cached.

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

As in :func:jaxamg.solve.

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.

A declared jaxamg.Pattern (from jaxamg.pattern); its colouring is attached. Exclusive with coloring.

Default for jaxamg.solve's labels ((count, labels), local rows in MPI mode).

Default for jaxamg.solve's label_sum.

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).

Floating dtype of the unknowns the operator will be applied to (default: the default float dtype). Discovery runs at this dtype, so a dtype-dependent operator gets the pattern it has at that precision.

class:ColoringInfo: the 5-tuple ``(rows, cols, colors, n_colors,

shape)for reattachment withwith_cache(..., coloring=...)``, carrying

the discovery dtype. Results are cached on the operator per dtype.

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 whose sparsity structure the plans are built from. Reuse the metadata only with the same CSR structure and entry order; values may change.

If True, the backward pass never transposes, so no transpose plan is built. Should match the is_symmetric passed to with_cache; the default (False) builds 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.