gvisor.dev/gvisor@v0.0.0-20240520182842-f9d4d51c7e0f/images/tpu/jax/Dockerfile.x86_64 (about)

     1  FROM python:3.8
     2  
     3  RUN pip install --upgrade jax
     4  RUN pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html