Enforcing jax[cuda] in a requirements.txt file

[GeoffNN]

Hi,
I'm writing a requirements.txt file for a package, which will be used by github.com/deepmind/xmanager. I want to force using cuda, and therefore want to encode this in my requirements.txt file. I've been struggling on this since simply writing the following line in the file does not work.

jax[cuda]>=VERSION -f https://storage.googleapis.com/jax-releases/jax_releases.html

[GeoffNN]

I've stumbled on a fix, by searching other github repos for similar situations. This commit gave me a working answer.

Adding the two following lines enforces the cuda install. I wanted to share this for others who may be similarly stumbled.

--find-links https://storage.googleapis.com/jax-releases/jax_releases.html
jax[cuda]>=version