Installation

aimz requires Python 3.11 or higher and is available via PyPI and conda-forge.

Quick Install

CPU (default)

pip install -U aimz

Installs aimz with CPU-only JAX — no GPU drivers required.

GPU (CUDA 13)

pip install -U "aimz[gpu]"

Installs aimz and jax[cuda13] for NVIDIA GPU acceleration.

Warning

GPU support is not available on Windows due to JAX limitations. Use Linux or WSL2 with a compatible NVIDIA GPU and CUDA drivers.

conda

conda install conda-forge::aimz

Installs the CPU version from conda-forge. GPU support via conda is not yet available — use pip for GPU installs.

Optional Extras

aimz ships several optional dependency groups that you can install with the pip install "aimz[<extra>]" syntax. Combine multiple extras with commas, e.g. pip install "aimz[mlflow,docs]".

mlflow

pip install -U "aimz[mlflow]"

Adds MLflow integration for experiment tracking and model logging.

docs

pip install -U "aimz[docs]"

Installs everything needed to build the documentation locally.

dev

pip install -U "aimz[dev]"

Development tools for linting, formatting, and testing.

Install from Source

For the latest unreleased changes, install directly from GitHub:

Latest (CPU)

pip install aimz@git+https://github.com/markean/aimz.git

Latest (GPU)

pip install "aimz[gpu]@git+https://github.com/markean/aimz.git"

Specific release

# Replace <tag> with the desired release tag
pip install aimz@git+https://github.com/markean/aimz.git@<tag>

Editable (development)

git clone https://github.com/markean/aimz.git
cd aimz
pip install -e ".[dev]"

Installs in editable mode with development dependencies.