Installation
============


From PyPI
---------

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

.. code-block:: bash

    pip install jaxion

For GPU support, use the following command instead:

.. code-block:: bash

    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:

.. code-block:: bash

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

Navigate to the project directory:

.. code-block:: bash

    cd jaxion

Install the package using pip (CPU version):

.. code-block:: bash

    pip install .

For GPU support, use the following command instead:

.. code-block:: bash

    pip install .[cuda12]

Verify the installation by running the test suite:

.. code-block:: bash

    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.