Development Guide

Notes for contributors working on JAX-AMG: setting up a development install, running the test suite, and the code-quality workflow.

Development Install

Follow the Installation Guide to set up CUDA, AmgX, and the build environment variables (CUDA_HOME, AMGX_ROOT, AMGX_BUILD). Then install in editable mode with the development and MPI extras:

pip install -e ".[all]"     # editable install + dev + mpi extras

This pulls in pytest, pytest-mpi, ruff, black, mypy, pre-commit, and the MPI dependencies. The equivalent pinned list lives in requirements-dev.txt.

Note

JAX-AMG links against a CUDA build of AmgX, so AmgX and CUDA must be on LD_LIBRARY_PATH at runtime (see Post-Installation Setup). Running with a bare path to the Python interpreter, instead of an activated environment, drops these variables and yields AMGX Error: Error initializing amgx core.

Running Tests

The suite uses pytest. Pin to a specific device with CUDA_VISIBLE_DEVICES.

Non-MPI tests

Tests that require MPI are marked @pytest.mark.mpi and are skipped by default, so a plain run exercises everything else on a single GPU:

CUDA_VISIBLE_DEVICES=0 python -m pytest tests/

MPI tests

The MPI tests (tests/test_mpi.py) must be launched under mpirun and enabled with the --only-mpi flag provided by pytest-mpi. Most need at least two ranks; assign a distinct GPU to each via CUDA_VISIBLE_DEVICES:

CUDA_VISIBLE_DEVICES=0,1 \
OMPI_MCA_opal_cuda_support=true MPI4JAX_USE_CUDA_MPI=1 \
mpirun -n 2 python -m pytest --only-mpi tests/

--only-mpi runs only the MPI-marked tests.
Omit the GPU-aware MPI variables (or set MPI4JAX_USE_CUDA_MPI=0) to stage communication through host memory when GPU-aware MPI is unavailable. See the MPI Guide and Environment Variables Reference for details.

JAX sharding tests

Run the sharding integration tests separately so every MPI process retains visibility of all participating GPUs:

CUDA_VISIBLE_DEVICES=0,1 \
mpirun -n 2 python -m pytest --only-mpi tests/test_sharding_mpi.py

Selecting this test module prevents the ordinary MPI test fixture from narrowing each rank to one visible CUDA device and enables early jax.distributed setup.

Tip

The MPI demos under demo/ (e.g. mpirun -n 4 python demo/mpi_autodiff.py) are a quick way to sanity-check distributed behavior end to end.

Code Quality

Formatting and linting are enforced with pre-commit (black, ruff with --fix, and mypy).

pre-commit install            # run automatically on each commit
pre-commit run --all-files    # run manually across the repo

Rebuilding the Native Extension

The solver is a C++ FFI extension. After editing any C++ source or header in jaxamg/, reinstall to rebuild it before testing; Python-only changes do not need a rebuild:

pip install -e . -v

MPI support in the extension is auto-enabled when AmgX was built with MPI; set JAXAMG_ENABLE_MPI=1 to force it on.

Building the Docs

The documentation is built with MkDocs from the sources in docs/ (configuration in mkdocs.yml). Preview it locally with:

mkdocs serve