Script to set up a working conda environment

[adrian-valente]

This repo is very cool, but due to numerous changes throughout the versions of jax, tensorflow and python since it was written it can be difficult to set up an environment where the functions can be profitably used. After some digging, I managed to pull out this script that works for me, albeit without CUDA (I didn't find the relevant versions of jax still available on their release repository).

If that can help people find this & build upon this script, I would be happy to leave this as an open PR and link from an issue, lmk if you know how to make it more usable / general.

Ps. We could also add the version of tensorflow itself if that can help, in my case 2.2.3 has been chosen by pip.

[google-cla[bot]]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.