feat: Support Python 3.12

[qnixsynapse]

This is the error I am getting:

jax_plugins/intel_extension_for_openxla/pjrt_plugin_xpu.so: undefined symbol: _ZNK4sycl3_V16detail16AccessorBaseHost25isMemoryObjectUsedByGraphEv

[fredlarochelle]

Yeah, it's not in an usable state yet, especially on Arc...

[rahulunair]

Any update on this?, I tried using 2024.0 and 2024.1 version of oneapi and with latest version of jax as well as the version prescribed in the readme. Both didnt work.

using 2024.0 version of oneapi and jax 0.4.25, here is the error I am getting on a Intel GPU Max system:

>>> import jax
>>> jax.local_devices()
INFO: Intel Extension for OpenXLA version: 0.3.0, commit: 9a484818
Jax plugin configuration error: Exception when calling jax_plugins.intel_extension_for_openxla.initialize()
Traceback (most recent call last):
  File "/home/sdp/.conda/envs/jax/lib/python3.11/site-packages/jax/_src/xla_bridge.py", line 482, in discover_pjrt_plugins
    plugin_module.initialize()
  File "/home/sdp/.conda/envs/jax/lib/python3.11/site-packages/jax_plugins/intel_extension_for_openxla/__init__.py", line 39, in initialize
    c_api = xb.register_plugin("xpu",
            ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/sdp/.conda/envs/jax/lib/python3.11/site-packages/jax/_src/xla_bridge.py", line 544, in register_plugin
    c_api = xla_client.load_pjrt_plugin_dynamically(plugin_name, library_path)  # type: ignore
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/sdp/.conda/envs/jax/lib/python3.11/site-packages/jaxlib/xla_client.py", line 155, in load_pjrt_plugin_dynamically
    return _xla.load_pjrt_plugin(plugin_name, library_path, c_api=None)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
jaxlib.xla_extension.XlaRuntimeError: INTERNAL: Failed to open /home/sdp/.conda/envs/jax/lib/python3.11/site-packages/jax_plugins/intel_extension_for_openxla/pjrt_plugin_xpu.so: /home/sdp/.conda/envs/jax/lib/python3.11/site-packages/jax_plugins/intel_extension_for_openxla/pjrt_plugin_xpu.so: undefined symbol: _ZNK4sycl3_V16detail16AccessorBaseHost25isMemoryObjectUsedByGraphEv
[CpuDevice(id=0)]
>>>

Using latest jax, 2024.1 oneapi and latest openxla I am getting a core dump.

[Zantares]

Any update on this?, I tried using 2024.0 and 2024.1 version of oneapi and with latest version of jax as well as the version prescribed in the readme. Both didnt work.

using 2024.0 version of oneapi and jax 0.4.25, here is the error I am getting on a Intel GPU Max system:

1. 2024.0 is not supported if you are using the Extension from pip. It's built by 2024.1 and will show undefined symbol error if works with 2024.0. Please rebuild Extension in your environment if you don't want to switch to 2024.1.
2. You mentioned that JAX v0.4.25 is used, so maybe you are suffering similar issue as installation issue #18. Each Extension version can only work with specific JAX version, such as 0.2.0 vs. 0.4.20. 0.3.0 vs. 0.4.24... Currently you can find the match info in the release notes: https://github.com/intel/intel-extension-for-openxla/releases .

[qnixsynapse]

I am currently using 2024.1 and it stopped detecting my GPU.

The output of jax.local_devices() shows [CpuDevice(id=0)]

[Zantares]

I am currently using 2024.1 and it stopped detecting my GPU.

The output of jax.local_devices() shows [CpuDevice(id=0)]

Do you meet the undefined symbol error even with 2024.1? If so, can you help to

• list the JAX version by pip list |egrep jax?
• list the Driver version by dpkg -l |egrep level-zero-gpu?

[qnixsynapse]

The undefined symbol is fixed when upgraded to 2024.1. However, it stopped detecting the pjrt plugin/GPU.

$ pip list |egrep jax
egrep: warning: egrep is obsolescent; using grep -E
jax                         0.4.25
jaxlib                      0.4.25

$ sudo pacman -Qs level-zero
local/intel-compute-runtime 24.13.29138.7-1
    Intel(R) Graphics Compute Runtime for oneAPI Level Zero and OpenCL(TM) Driver
local/level-zero-loader 1.16.15-1
    API for accessing low level interfaces in oneAPI platform devices (loader)

[Zantares]

@qnixsynapse JAX v0.4.25 doesn't match released Extension v0.3.0, so please try to use any of the following 2 solutions first:

1. Downgrade JAX version to v0.4.24
2. Rebuild the newest Extension main branch by yourself, it has upgraded to support JAX v0.4.25

More version-matching info can be found in the release notes: https://github.com/intel/intel-extension-for-openxla/releases. We will add a table in the home page to tell the matching info later.

[qnixsynapse]

@Zantares Same issue with v0.4.24:

Edit: It seems a python version mismatch, reason why pip is pulling a pre release old pywheel of openxla plugin. Arch Linux has python version: 3.12.3, the new wheels are built with 3.11.

[Zantares]

@Zantares Same issue with v0.4.24:

Edit: It seems a python version mismatch, reason why pip is pulling a pre release old pywheel of openxla plugin. Arch Linux has python version: 3.12.3, the new wheels are built with 3.11.

Thanks for the trial and the new info. We will check the error with Python 3.12 on Arc.

BTW, isn't Extension v0.3.0 installed but an old version if it's mismatched?

[qnixsynapse]

Seems like it. Pip is refusing to install the 0.3.0 because it is trying to search for a wheel built for Python 3.12

I guess the last one is the one getting installed.

Also, this:

[Zantares]

That's it. We will discuss it internally first to see if we can release more Python packages (means more test processes). Right now you can try to build it by yourself if Python 3.12 is hard requested.