JAX+Pallas installation

[cpcpython]

Hi list

When I tried to import pallas it failed (I have default Jax, version: '0.4.13' in TPU v2-8 machine ), and it print:

>>> from jax.experimental import pallas as pl
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ImportError: cannot import name 'pallas' from 'jax.experimental' (/home/gpkmohan_mbcet/.local/lib/python3.8/site-packages/jax/experimental/__init__.py)
>>> from jax.experimental.pallas import tpu as pltpu
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ModuleNotFoundError: No module named 'jax.experimental.pallas'

I tried:

pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html --upgrade --force-reinstall jax==0.4.35

But It failed.

Can anyone help us for this?

Thanks in advance

Mohan

[mattjj]

Thanks for the question!

But It failed.

Can you give more detail about this? What was the error?

/home/gpkmohan_mbcet/.local/lib/python3.8/site-packages/jax/...

The issue may be that Python 3.8 is too old, and you should update Python.

[jakevdp]

JAX v0.4.13 was the last version to support Python 3.8 (see changelog), so this is almost certainly the issue.

[cpcpython]

Thanks for the suggestionl for Python 3.10 its ok Now!!!

Python 3.10.6 (main, Mar 10 2023, 10:55:28) [GCC 11.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import pallas
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ModuleNotFoundError: No module named 'pallas'
>>> import jax
>>> jax.__version__
'0.4.35'
>>> from jax.experimental import pallas as pl
>>> 

[cpcpython]

Closing