Pip install gives error on TPU clustere

[abhinavg4]

The following setup commands:

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

Given the following error:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/ubuntu/.local/lib/python3.10/site-packages/levanter/__init__.py", line 4, in <module>
    import levanter.distributed as distributed
  File "/home/ubuntu/.local/lib/python3.10/site-packages/levanter/distributed.py", line 12, in <module>
    from jax._src.clusters import SlurmCluster, TpuCluster
ImportError: cannot import name 'TpuCluster' from 'jax._src.clusters' (/home/ubuntu/.local/lib/python3.10/site-packages/jax/_src/clusters/__init__.py)

Building from source works fine

[dlwh]

the latest levanter release is super old. you need to use the git install