Incompatibility with jaxlib 0.3.7

The newest version of jax seems to require jaxlib v0.3.7, which breaks the trainer script:
```bash
$ ./run_pretrain.sh
2022-04-16 23:34:03.151271: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
Traceback (most recent call last):
File "/usr/lib/python3.8/runpy.py", line 185, in _run_module_as_main
mod_name, mod_spec, code = _get_module_details(mod_name, _Error)
File "/usr/lib/python3.8/runpy.py", line 111, in _get_module_details
__import__(pkg_name)
File "/data/t5x/t5x/__init__.py", line 17, in
import t5x.adafactor
File "/data/t5x/t5x/adafactor.py", line 63, in
from t5x import utils
File "/data/t5x/t5x/utils.py", line 41, in
from t5x import checkpoints
File "/data/t5x/t5x/checkpoints.py", line 51, in
from t5x import optimizers
File "/data/t5x/t5x/optimizers.py", line 36, in
import optax
File "/data/venvt5/lib/python3.8/site-packages/optax/__init__.py", line 17, in
from optax import experimental
File "/data/venvt5/lib/python3.8/site-packages/optax/experimental/__init__.py", line 20, in
from optax._src.experimental.complex_valued import split_real_and_imaginary
File "/data/venvt5/lib/python3.8/site-packages/optax/_src/experimental/complex_valued.py", line 32, in
import chex
File "/data/venvt5/lib/python3.8/site-packages/chex/__init__.py", line 17, in
from chex._src.asserts import assert_axis_dimension
File "/data/venvt5/lib/python3.8/site-packages/chex/_src/asserts.py", line 26, in
from chex._src import asserts_internal as _ai
File "/data/venvt5/lib/python3.8/site-packages/chex/_src/asserts_internal.py", line 32, in
from chex._src import pytypes
File "/data/venvt5/lib/python3.8/site-packages/chex/_src/pytypes.py", line 40, in
CpuDevice = jax.lib.xla_extension.CpuDevice
AttributeError: module 'jaxlib.xla_extension' has no attribute 'CpuDevice'
```
Forcing the install of `"jax[tpu]<0.3.7"` works for now.
```bash
pip install -U "jax[tpu]<0.3.7" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
```

Start by reviewing the project's Python dependency declarations and run_pretrain.sh, then reproduce the import failure with jaxlib 0.3.7. Check the installed jax, optax, and chex versions against the reported traceback. Done means the trainer script starts without the CpuDevice error while preserving a documented compatible dependency set.

Written by the indexing model from the issue text.