--find-links https://pypi.org/simple/
--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
jaxlib==0.3.22+cuda11.cudnn82
jax==0.3.22
trax==1.4.1