Installation

From PyPI

To install the latest release version of the Jaxion package, run the following command:

pip install jaxion

For GPU support, use the following command instead:

pip install jaxion[cuda12]

Note

For now, to build with GPU support, use the build from source method below. This will be simplified in future releases.

Build from Source

Check out the repository:

git clone git@github.com:JaxionProject/jaxion.git

Navigate to the project directory:

cd jaxion

Install the package using pip (CPU version):

pip install .

For GPU support, use the following command instead:

pip install .[cuda12]

Verify the installation by running the test suite:

pytest

Build Prerequisites

Jaxion is tested primarily on Linux with Python 3.11, 3.12, and 3.13. Newer Python releases may not be supported until the dependency stack publishes compatible wheels.

Jaxion depends on jaxdecomp. When a jaxdecomp wheel is not available for your platform, pip builds it from source during installation. This source build requires a C++ toolchain, cmake, and ninja to be available on the build path. On macOS arm64, jaxdecomp may currently be built from source because no wheel is available. CUDA is not required for the CPU install; if a CUDA compiler is not found, jaxdecomp falls back to its pure-JAX backend.