Installation

First, make sure that you have a working installation of JAX (as explained here)

Then, use the following command will install augmax from PyPI:

pip install augmax

In case you want to use the latest version from github directly, please run

pip install git+https://github.com/khdlr/augmax