Installation

Requirements
Installing via Conda (Recommended)
Installing from source (pip only)
GPU Memory Allocation

Requirements

jaxpe relies on JAX and its ecosystem for automatic differentiation, GPU acceleration, and vectorization.

JAX
flowjax — Rational-quadratic-spline coupling flows
standard scientific stack: numpy, scipy, matplotlib

Installing via Conda (Recommended)

For the most robust installation—especially when dealing with complex dependencies like lalsuite and GPU acceleration—we recommend using the provided Conda environment files. These will automatically configure lalsuite through conda-forge and install jaxpe with its optional dependencies (dev, gwdata, surrogate, and esigma) using pip.

The esigma extra (diffrax, numba, numbalsoda) is needed for the ESIGMAInspiral waveform model (--waveform esigma); it also requires a separate editable install of esigmapy-dev, which conda.yml/conda-gpu.yml currently reference by a local path — edit that line to point at your own esigmapy-dev checkout before running conda env create.

For CPU:

conda env create -f conda.yml
conda activate jaxpe

For GPU (NVIDIA): This environment specifically prioritizes the official jax[cuda12] binaries to ensure hardware acceleration is properly linked before installing jaxpe.

conda env create -f conda-gpu.yml
conda activate jaxpe-gpu

Installing from source (pip only)

Alternatively, the package can be installed manually via pip:

git clone https://github.com/jaxpe/jaxpe.git
cd jaxpe
pip install -e .

Ensure that you have installed the correct JAX version with GPU support (if a GPU is available) by following the JAX installation instructions.

GPU Memory Allocation

By default, JAX pre-allocates 90% of the available GPU memory. When running heavily vmapped applications or using large batch sizes (like in GW PE), this might need tuning.

If you are running on smaller GPUs or alongside other workloads, you may need to restrict JAX’s memory allocation:

export XLA_PYTHON_CLIENT_MEM_FRACTION=0.5

This prevents on-demand allocation fragmentation which can sometimes cause issues mid-run.