Installation¶

From PyPI (soon)

pip install apax

From GitHub¶

If you would like to have a pre-release version, you can install Apax from GitHub directly

pip install git+https://github.com/apax-hub/apax.git

As a developer, you first need to install Poetry. You can obtain it by running

curl -sSL https://install.python-poetry.org | python3 -

You can then clone and install the project.

git clone https://github.com/apax-hub/apax.git <dest_dir>
cd <dest_dir>
poetry install

Note that all of the above only install the CPU version. If you want to enable GPU support, please overwrite the jaxlib version:

CUDA 12:

pip install -U "jax[cuda12]"