Unable to install t5x on cloud TPU

[hxssgaa]

Hi,

I have been following the installation instruction here.

However I encounted the following error when installing the t5x library on cloud TPU v4-8 (tpu version: tpu-ubuntu2204-base):

ERROR: Could not find a version that satisfies the requirement flax (unavailable) (from t5x[tpu]) (from versions: 0.1.0rc1, 0.1.0rc2, 0.1.0, 0.2.0, 0.2.1, 0.2.2, 0.3.0, 0.3.1, 0.3.2, 0.3.3, 0.3.4, 0.3.5, 0.3.6, 0.4.0, 0.4.1, 0.4.2, 0.5.0, 0.5.1, 0.5.2, 0.5.3, 0.6.0, 0.6.1, 0.6.2, 0.6.3, 0.6.4, 0.6.5, 0.6.6, 0.6.7, 0.6.8, 0.6.9, 0.6.10, 0.6.11, 0.7.0, 0.7.1, 0.7.2, 0.7.3, 0.7.4)
ERROR: No matching distribution found for flax (unavailable)

May I know how to resolve the issue?

[rehaidib]

I have the same issue

[tomlimi]

me too

[DanielWFrancis]

Me too

[stefan-it]

Workaround for this:

Create a v3 or v4 TPU VM (I have tested this on a v3-32 and v4-32 TPU Pod) using the tpu-ubuntu2204-base runtime version, e.g. with resource manager it can be created like:

$ gcloud alpha compute tpus queued-resources create byt5-experiments \
  --node-id byt5 --project <project-name> \
  --zone us-central2-b \
  --accelerator-type v4-32 \
  --runtime-version tpu-ubuntu2204-base

On each worker, perform the following steps:

Install jax on each worker:

$ pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

Install latest flax (here: 8360b7c)

$ pip3 install --upgrade git+https://github.com/google/flax.git

Install latest T5X version (here: b642f30)

$ git clone https://github.com/google-research/t5x.git
$ sed -i 's!flax @ git+https://github.com/google/flax#egg=flax!flax!g' t5x/setup.py
$ cd t5x && pip3 install -e
$ cd ..

Then T5X can be used :)