Notebook implementing Bayesian Model Averaging with numpyro.
Run the environment.yml file by running the following command on the main repo directory:
conda env create
The installation works for conda==4.12.0. This will install all packages needed to run the code on a CPU with jupyter.
If you want to run this code with a CUDA GPU, you will need to download the appropriate jaxlib==0.4.13 version. For example, for my GPU running on CUDA==12.3, I would run:
pip install jaxlib==0.4.13+cuda12.cudnn89
The key to using this code directly would be to retain the jax and jaxlib versions.