Installation

The standard installation installs the CPU version of Apax. To enable GPU support (available only on Linux), install Apax with the extra CUDA dependencies. See the Jax installation instructions for more details.

From PyPI

CPU:

pip install apax

GPU:

pip install "apax[cuda]"

From GitHub

For a pre-release version, install Apax directly from GitHub.

CPU:

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

GPU:

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

For Developers

To set up a development environment, first install uv.

pip install uv

Then clone the project from GitHub,

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

and install it.

CPU:

uv sync --all-extras --no-extra cuda

GPU:

uv sync --extra cuda

Extra Dependencies

If you want to use Apax in the IPSuite framework and use the predefined apax.nodes, you can install the extra dependencies for IPSuite:

pip install "apax[ipsuite]"

Additionally, you have the option to install the extra dependencies for MLFlow:

pip install "apax[mlflow]"

For hyperparameter optimization, optuna is required, which can be installed as

pip install "apax[optuna]"