Installation

Step 1 - Install JAX

This step is necessary only if you need to run codes on GPU.

pip install -U "jax[cuda12]"

Check JAX installation for other installation options.

Step 2 - Install Quantax

git clone -b dev https://github.com/ChenAo-Phys/quantax.git
pip install ./quantax