Skip to content

Installation

Quick Start

uv add jaxlatt

pip install jaxlatt

Requirements

  • Python 3.9 or later
  • JAX 0.4.0 or later
  • NumPy, Matplotlib

Installation Options

Standard Installation

For CPU-only usage:

uv add jaxlatt

pip install jaxlatt

GPU Support

To enable GPU acceleration with CUDA:

uv add "jax[cuda12]"
uv add jaxlatt

pip install -U "jax[cuda12]"
pip install jaxlatt

See the JAX installation guide for platform-specific GPU setup.

Development Installation

To install from source with development dependencies:

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

uv pip install -e ".[dev]"

pip install -e ".[dev]"

Verify Installation

Test your installation:

import jax
import jaxlatt

print(f"JAX backend: {jax.default_backend()}")
print(f"JaxLatt version: {jaxlatt.__version__}")

Troubleshooting

JAX not using GPU?

  • Verify CUDA installation: nvidia-smi
  • Check JAX sees GPU: jax.devices()
  • Ensure JAX GPU variant is installed

Import errors?

  • Upgrade pip: pip install --upgrade pip
  • Install in a clean environment: python -m venv venv && source venv/bin/activate

For more help, see GitHub Issues.