Installation

Prerequisites

archetypax requires Python 3.8 or later. The package leverages JAX for GPU acceleration, which necessitates compatible hardware for optimal performance.

Standard Installation

You can install archetypax directly from PyPI using pip:

pip install archetypax

This will install the core package with all required dependencies.

Development Installation

For development purposes, we recommend installing from source with development dependencies:

git clone https://github.com/yourusername/archetypax.git
cd archetypax
pip install -e ".[dev,docs,examples]"

This installs the package in editable mode with additional dependencies for development, documentation, and examples.

GPU Acceleration

To leverage GPU acceleration, ensure you have installed the appropriate JAX version for your hardware:

# For CUDA-compatible NVIDIA GPUs
pip install --upgrade "jax[cuda12]"

# For TPUs
pip install --upgrade "jax[tpu]"

Verification

To verify your installation, you can run:

import archetypax
print(archetypax.__version__)

# Check if JAX can detect your accelerator
import jax
print(jax.devices())