jax[tpu]

[ma7dev]

Hi,

I have been trying to install jax[tpu] using poetry but I couldn't. I have been using this command to install the jax version that works with tpu using pip:

python -m pip install --upgrade "jax[tpu]==0.3.10" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

However, I couldn't find a way to have it work with poetry. Do you have any ideas? Thanks!

[jorenham]

If you configure your pyproject.toml like done here, it should work when you poetry add jax[tpu].
https://github.com/jorenham/jax_pep503/blob/master/example_project/pyproject.toml

[ma7dev]

I have added this and it worked out! Thanks!

jaxlib = { extras = ["tpu"], version = "0.3.10" }