Installation

Requirements

splax needs an NVIDIA GPU and a CUDA-enabled JAX. The jax[cuda] wheel is pulled in as a dependency, so no system CUDA toolchain is required. Python 3.11 is required.

Install with pip

uv pip install "git+https://github.com/learnsyslab/splax"

Verify the import:

python -c "import splax; print(splax.__all__)"

Developer setup

pixi installs splax editable with the developer tooling into a managed environment.

pixipixi + docs

git clone https://github.com/learnsyslab/splax.git
cd splax
pixi shell

git clone https://github.com/learnsyslab/splax.git
cd splax
pixi shell -e docs

Use the pixi environments above for tooling and docs tasks.

Building the documentation

The docs are built with ProperDocs and the Material theme. Set JAX_PLATFORMS=cpu so the API reference imports splax without a GPU:

JAX_PLATFORMS=cpu pixi run -e docs docs-build