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 --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

CUDA 11:

pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html