Installation
ReLax
TL;DR: Install ReLax
via the Python Package Index:
pip install jax-relax
ReLax
only support installation on Linux and macOS. This is because ReLax
is built on top of Jax, which currently does not support official Windows wheels.
Windows users should run ReLax
on CPU and GPU on Windows Subsystem for Linux.
Installing ReLax
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.8 -y
conda activate relax
ReLax
is built on top of Jax. You should also check the official installation guide from the Jax team.
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:
git clone https://github.com/BirkhoffG/relax.git
pip install -e relax
Running on GPU or TPU
We do not run continuous integration (CI) for GPU and TPU environments. So ReLax
might fail unexpectedly when running on GPU/TPU. If you encounter issues when running on GPU/TPU, please report to us.
If you wish to run relax on GPU or TPU, please first install this library via pip install jax-relax
.
Then, you should install the right GPU or TPU version of Jax by following steps in the install guidelines.
Contributor of relax
See nbdev installation.