From fb487fbb2fa1b2596d38633b1fedb2f2348bb711 Mon Sep 17 00:00:00 2001 From: Lucas Manning Date: Wed, 18 Oct 2023 09:54:25 -0700 Subject: [PATCH] Create a TPU hello world image for use in gVisor tests. PiperOrigin-RevId: 574506997 --- images/tpu/jax/Dockerfile.x86_64 | 4 ++++ 1 file changed, 4 insertions(+) create mode 100644 images/tpu/jax/Dockerfile.x86_64 diff --git a/images/tpu/jax/Dockerfile.x86_64 b/images/tpu/jax/Dockerfile.x86_64 new file mode 100644 index 000000000..d1e088b1c --- /dev/null +++ b/images/tpu/jax/Dockerfile.x86_64 @@ -0,0 +1,4 @@ +FROM python:3.8 + +RUN pip install --upgrade jax +RUN pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html