Installation

A tutorial on installing ReLax
Tip

TL;DR: For most users, install ReLax via the Python Package Index:

pip install -U jax-relax
Warning

ReLax supports various platforms, thanks to efforts from developers of JAX. See the compatible platforms of JAX.

Installing ReLax

This section assumes that you are an end-user of ReLax, e.g., you only want to use this library for your own developement without modifying the ReLax codebase.

ReLax is built on top of JAX. You should also check the official installation guide from the Jax team.

Prerequisite: Set up your python environment

We suggest to create a new environment when using ReLax.

If you are using conda, you can create a new environment by:

conda create -n relax python=3.9 -y
conda activate relax

Running on CPU

If you only need to run relax on CPU, you can simply install via pip or clone the GitHub project.

Installation via PyPI:

pip install --upgrade pip
pip install --upgrade jax-relax

Editable Install

If you wish to install ReLax from source, run:

# Clone the ReLax Repository
git clone https://github.com/BirkhoffG/jax-relax.git
cd jax-relax
# Clone the submodule
git submodule update --init --recursive
pip install -e .

Running on GPU or TPU

If you wish to run ReLax on GPU or TPU, please first install this library via pypi. Next, you should install the right GPU or TPU version of JAX by following steps in the install guidelines.

For example, if you want to install a GPU version, you should run

pip install jax-relax

Next, install the GPU version of jax:

pip install -U "jax[cuda12]"
Warning

We do not run continuous integration (CI) for GPU and TPU environments. If you encounter issues when running on GPU/TPU, please report to us.

If you are a Contributor of ReLax…

You will need to install additional packages if you want to fork and make changes to the library.

You should install the required packages via

pip install "jax-relax[dev]"

Next, follow the nbdev installation for further instructions. Check out the contribution guidance for more details.