Install

JAX-AM supports Linux and macOS. Create a new conda environment and run

git clone https://github.com/tianjuxue/jax-am.git
cd jax-am
pip install .

Note: JAX-AM depends on petsc4py. We have found difficulty installing petsc4py with pip on certain platforms. Installing petsc4py with conda is therefore recommended (see here).