installation issue

[jinz2014]

After following the install instructions, I got the following error. Thanks

pip install jax==0.4.20 jaxlib==0.4.20 pip3 install --upgrade intel-extension-for-openxla

>>> import jax
>>> import jax.numpy as jnp
>>> import jax
>>> print("jax.local_devices(): ", jax.local_devices())
Platform 'xpu' is experimental and not all JAX functionality may be correctly supported!
Traceback (most recent call last):
  File "/home/users/test/.local/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 623, in backends
    backend = _init_backend(platform)
  File "/home/users/test/.local/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 734, in _init_backend
    backend = registration.factory()
  File "/home/users/test/.local/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 470, in factory
    xla_client.initialize_pjrt_plugin(plugin_name)
  File "/home/users/test/.local/lib/python3.10/site-packages/jaxlib/xla_client.py", line 159, in initialize_pjrt_plugin
    _xla.initialize_pjrt_plugin(plugin_name)
jaxlib.xla_extension.XlaRuntimeError: INVALID_ARGUMENT: Mismatched PJRT plugin PJRT API version (0.2) and framework PJRT API version 0.38).

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/users/test/.local/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 877, in local_devices
    process_index = get_backend(backend).process_index()
  File "/home/users/test/.local/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 779, in get_backend
    return _get_backend_uncached(platform)
  File "/home/users/test/.local/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 759, in _get_backend_uncached
    bs = backends()
  File "/home/users/test/.local/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 639, in backends
    raise RuntimeError(err_msg)
RuntimeError: Unable to initialize backend 'xpu': INVALID_ARGUMENT: Mismatched PJRT plugin PJRT API version (0.2) and framework PJRT API version 0.38). (you may need to uninstall the failing plugin package, or set JAX_PLATFORMS=cpu to skip this backend.)

[Zantares]

The most likely reason is that you are using an Extension that does not match the JAX version. Each Extension version can only work with specific JAX version, such as 0.1.0 vs. 0.4.13, 0.2.0 vs. 0.4.20... Currently you can find the match info in release notes: https://github.com/intel/intel-extension-for-openxla/releases

We will maintain a clear table in the future in case you miss it, thanks.

[Zantares]

This question has been inactive for more than 2 weeks, will close it soon if no more feedback, thanks!