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](https://jax.readthedocs.io/en/latest/installation.html) for
other installation options.

Step 2 - Install Quantax
---------------------------

.. code-block:: python

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


Supported platforms
-------------------
CPU, NVIDIA GPU