- changed version from 3.1.1 to 1.5.0 - removed pip install jax - added mamba install -c conda-forge jax