Paper reference: arxiv:2402.11014
Author: Matija Medvidović
Make sure you have JAX and other dependencies installed. You can do this in your terminal by navigating to the repo root folder and running:
virtualenv venv
source venv/bin/activate
pip install .Do this outside of the notebook environment and then make sure to select the created virtual environment as the kernel in any notebooks you run.