Rebase 0.4.38 (#569)

Author: akhilgoe

.github/CODEOWNERS

@@ -1,2 +1,2 @@
 
-# * @agramesh1 @vsanghavi @mdfaijul @akhilgoe @ShengYang1 @bhavani-subramanian @ashraf-bhuiyan
+* @agramesh1 @vsanghavi @mdfaijul @akhilgoe @Solaryee @bhavani-subramanian @ashraf-bhuiyan

CONTRIBUTING.md

@@ -2,7 +2,7 @@
 
 ### License
 
-<PROJECT NAME> is licensed under the terms in [LICENSE]<link to license file in repo>. By contributing to the project, you agree to the license and copyright terms therein and release your contribution under these terms.
+Intel® Extension for OpenXLA is licensed under the terms in [LICENSE](LICENSE.txt). By contributing to the project, you agree to the license and copyright terms therein and release your contribution under these terms.
 
 ### Sign your work
 

README.md

@@ -7,7 +7,7 @@
 
 The [OpenXLA](https://github.com/openxla/xla) Project brings together a community of developers and leading AI/ML teams to accelerate ML and address infrastructure fragmentation across ML frameworks and hardware.
 
-Intel® Extension for OpenXLA includes PJRT plugin implementation, which seamlessly runs JAX models on Intel GPU. The PJRT API simplified the integration, which allowed the Intel GPU plugin to be developed separately and quickly integrated into JAX. This same PJRT implementation also enables initial Intel GPU support for TensorFlow and PyTorch models with XLA acceleration. Refer to [OpenXLA PJRT Plugin RFC](https://github.com/openxla/community/blob/main/rfcs/20230123-pjrt-plugin.md) for more details.
+Intel® Extension for OpenXLA includes PJRT plugin implementation, which seamlessly runs JAX models on Intel GPU. The PJRT API simplifies the integration, enabling the Intel GPU plugin to be developed separately and allowing smooth integration with JAX. This same PJRT implementation also enables initial Intel GPU support for TensorFlow and PyTorch models with XLA acceleration. Refer to [OpenXLA PJRT Plugin RFC](https://github.com/openxla/community/blob/main/rfcs/20230123-pjrt-plugin.md) for more details.
 
 This guide introduces the overview of OpenXLA high level integration structure and demonstrates how to build Intel® Extension for OpenXLA and run JAX example with OpenXLA on Intel GPU. JAX is the first supported front-end.
 
@@ -26,59 +26,48 @@ This guide introduces the overview of OpenXLA high level integration structure a
 
 Verified Hardware Platforms:
 
-* Intel® Data Center GPU Max Series, Driver Version: [LTS release 2350.125](https://dgpu-docs.intel.com/releases/LTS-release-notes.html)
+* Intel® Data Center GPU Max Series, Driver Version: [LTS release 2350.136](https://dgpu-docs.intel.com/releases/LTS-release-notes.html)
 
-* Intel® Data Center GPU Flex Series, Driver Version: [LTS release 2350.125](https://dgpu-docs.intel.com/driver/installation.html)
+* Intel® Data Center GPU Flex Series, Driver Version: [LTS release 2350.136](https://dgpu-docs.intel.com/driver/installation.html)
 
 ### Software Requirements
 
 * Ubuntu 22.04 (64-bit)
   * Intel® Data Center GPU Flex Series
 * Ubuntu 22.04, SUSE Linux Enterprise Server(SLES) 15 SP4
   * Intel® Data Center GPU Max Series
-* [Intel® oneAPI Base Toolkit 2025.0](https://www.intel.com/content/www/us/en/developer/articles/release-notes/intel-oneapi-toolkit-release-notes.html)
-* Jax/Jaxlib 0.4.30
-* Python 3.9-3.12
+* [Intel® Deep Learning Essentials 2025.1](https://www.intel.com/content/www/us/en/developer/tools/oneapi/base-toolkit-download.html?packages=dl-essentials&dl-lin=offline&dl-essentials-os=linux)
+* Jax/Jaxlib 0.4.38
+* Python 3.9-3.13
 * pip 19.0 or later (requires manylinux2014 support)
 
-**NOTE: Since JAX has its own [platform limitation](https://jax.readthedocs.io/en/latest/installation.html#supported-platforms) (Ubuntu 20.04 or later), real software requirements is restricted when works with JAX.**
+**NOTE: Since JAX has its own [platform limitation](https://jax.readthedocs.io/en/latest/installation.html#supported-platforms) (Ubuntu 20.04 or later), real software requirements are restricted when working with JAX.**
 
 ### Install Intel GPU Drivers
 
 |OS|Intel GPU|Install Intel GPU Driver|
 |-|-|-|
-|Ubuntu 22.04 |Intel® Data Center GPU Flex Series|  Refer to the [Installation Guides](https://dgpu-docs.intel.com/driver/installation.html) for latest driver installation. If install the verified Intel® Data Center GPU Max Series/Intel® Data Center GPU Flex Series [LTS release 2350.125](https://dgpu-docs.intel.com/driver/installation.html), please append the specific version after components, such as `sudo apt-get install intel-opencl-icd==24.45.31740.10-1057~22.04`|
-|Ubuntu 22.04, SLES 15 SP4|Intel® Data Center GPU Max Series|  Refer to the [Installation Guides](https://dgpu-docs.intel.com/driver/installation.html) for latest driver installation. If install the verified Intel® Data Center GPU Max Series/Intel® Data Center GPU Flex Series [LTS release 2350.125](https://dgpu-docs.intel.com/releases/LTS-release-notes.html), please append the specific version after components, such as `sudo apt-get install intel-opencl-icd==24.45.31740.10-1057~22.04`|
+|Ubuntu 22.04 |Intel® Data Center GPU Flex Series|  Refer to the [Installation Guides](https://dgpu-docs.intel.com/driver/installation.html) for latest driver installation. If installing the verified Intel® Data Center GPU Max Series/Intel® Data Center GPU Flex Series [LTS release 2350.136](https://dgpu-docs.intel.com/driver/installation.html), please append the specific version after components, such as `sudo apt-get install intel-opencl-icd==24.45.31740.10-1057~22.04`|
+|Ubuntu 22.04, SLES 15 SP4|Intel® Data Center GPU Max Series|  Refer to the [Installation Guides](https://dgpu-docs.intel.com/driver/installation.html) for latest driver installation. If installing the verified Intel® Data Center GPU Max Series/Intel® Data Center GPU Flex Series [LTS release 2350.136](https://dgpu-docs.intel.com/releases/LTS-release-notes.html), please append the specific version after components, such as `sudo apt-get install intel-opencl-icd==24.45.31740.10-1057~22.04`|
 
-### Install oneAPI Base Toolkit Packages
+### Install Intel® Deep Learning Essentials Packages
 
-Need to install components of Intel® oneAPI Base Toolkit:
+Need to install components of Intel® Deep Learning Essentials:
 
 * Intel® oneAPI DPC++ Compiler
 * Intel® oneAPI Math Kernel Library (oneMKL)
 
 ```bash
-$ wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/96aa5993-5b22-4a9b-91ab-da679f422594/intel-oneapi-base-toolkit-2025.0.0.885_offline.sh
+$ wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/e7705a6d-954d-465c-a5bc-4f820e2e4e90/intel-deep-learning-essentials-2025.0.2.9_offline.sh
 # 2 components are necessary: DPC++/C++ Compiler and oneMKL
-sudo sh intel-oneapi-base-toolkit-2025.0.0.885_offline.sh
+sudo sh ./intel-deep-learning-essentials-2025.0.2.9_offline.sh -a --silent --eula accept
 
 # Source OneAPI env
-source /opt/intel/oneapi/compiler/2025.0/env/vars.sh
-source /opt/intel/oneapi/mkl/2025.0/env/vars.sh
+source /opt/intel/oneapi/compiler/2025.1/env/vars.sh
+source /opt/intel/oneapi/mkl/2025.1/env/vars.sh
 export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/intel/oneapi/umf/latest/lib
 ```
 
-**Backup**: Recommend to rollback to **Toolkit 2024.1** if meet performance issue. See [Release Notes](https://github.com/intel/intel-extension-for-openxla/releases) for more details.
-```bash
-wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/fdc7a2bc-b7a8-47eb-8876-de6201297144/l_BaseKit_p_2024.1.0.596.sh
-# 2 components are necessary: DPC++/C++ Compiler and oneMKL
-sudo sh l_BaseKit_p_2024.1.0.596.sh
-
-# Source OneAPI env
-source /opt/intel/oneapi/compiler/2024.1/env/vars.sh
-source /opt/intel/oneapi/mkl/2024.1/env/vars.sh
-```
-
 ### Install Jax and Jaxlib
 
 ```bash
@@ -89,6 +78,7 @@ Please refer to [test/requirements.txt](test/requirements.txt) for the version d
 The following table tracks intel-extension-for-openxla versions and compatible versions of `jax` and `jaxlib`. The compatibility between `jax` and `jaxlib` is maintained through JAX. This version restriction will be relaxed over time as the plugin API matures.
 |**intel-extension-for-openxla**|**jaxlib**|**jax**|
 |:-:|:-:|:-:|
+| 0.6.0 | 0.4.38 | 0.4.38 |
 | 0.5.0 | 0.4.30 | >= 0.4.30, <= 0.4.31|
 | 0.4.0 | 0.4.26 | >= 0.4.26, <= 0.4.27|
 | 0.3.0 | 0.4.24 | >= 0.4.24, <= 0.4.27|
@@ -112,7 +102,7 @@ git clone https://github.com/intel/intel-extension-for-openxla.git
 ./configure        # Choose Yes for all.
 bazel build //xla/tools/pip_package:build_pip_package
 ./bazel-bin/xla/tools/pip_package/build_pip_package ./
-pip install intel_extension_for_openxla-0.5.0-cp39-cp39-linux_x86_64.whl
+pip install intel_extension_for_openxla-0.6.0-cp312-cp312-linux_x86_64.whl
 ```
 
 **Aditional Build Option**:
@@ -128,7 +118,7 @@ bazel build --override_repository=xla=/path/to/xla //xla/tools/pip_package:build
 By default, bazel will automatically search for the required libraries on your system. This eliminates the need for manual configuration in most cases. For more advanced use cases, you can specify a custom location for the libraries using environment variables:
 
 ```bash
-export MKL_INSTALL_PATH=/opt/intel/oneapi/mkl/2025.0
+export MKL_INSTALL_PATH=/opt/intel/oneapi/mkl/2025.1
 export L0_INSTALL_PATH=/usr
 bazel build //xla/tools/pip_package:build_pip_package
 ```
@@ -162,7 +152,7 @@ print(lax_conv())
 ### Reference result
 
 ```bash
-jax.local_devices():  [xpu(id=0), xpu(id=1)]
+jax.local_devices():  [sycl(id=0), sycl(id=1)]
 [[[[2.0449753 2.093208  2.1844783 1.9769732 1.5857391 1.6942389]
    [1.9218378 2.2862523 2.1549542 1.8367321 1.3978379 1.3860377]
    [1.9456574 2.062028  2.0365305 1.901286  1.5255247 1.1421617]
@@ -189,7 +179,7 @@ jax.local_devices():  [xpu(id=0), xpu(id=1)]
 2. If there is an error 'version GLIBCXX_3.4.30' not found, upgrade libstdc++ to the latest, for example for conda
 
     ```bash
-    conda install libstdcxx-ng==12.2.0 -c conda-forge
+    conda install libstdcxx-ng -c conda-forge -y
     ```
 
 3. If there is an error '/usr/bin/ld: cannot find -lstdc++: No such file or directory' during source build under Ubuntu 22.04, check the selected GCC-toolchain path and the installed libstdc++.so library path, then create symbolic link of the selected GCC-toolchain path to the libstdc++.so path, for example:

WORKSPACE

@@ -3,7 +3,7 @@ workspace(name = "intel_extension_for_openxla")
 load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 load("//third_party:version_check.bzl", "check_bazel_version_at_least")
 
-check_bazel_version_at_least("5.3.0")
+check_bazel_version_at_least("6.5.0")
 
 # To update XLA to a new revision,
 # a) update URL and strip_prefix to the new git commit hash
@@ -14,10 +14,10 @@ http_archive(
     name = "xla",
     patch_args = ["-p1"],
     patches = ["//third_party:openxla.patch"],
-    sha256 = "083c7281a629647ab2cc32f054afec74893c33e75328783b8085c818f48235ff",
-    strip_prefix = "xla-79fd5733f99b3c0948d7202bc1bbe1ee3980da5c",
+    sha256 = "0870fcd86678cae31c56cfc57018f52ceec8e4691472af62c847ade746a0eb13",
+    strip_prefix = "xla-20a482597b7dd3067b26ca382b88084ee5a21cf7",
     urls = [
-        "https://github.com/openxla/xla/archive/79fd5733f99b3c0948d7202bc1bbe1ee3980da5c.tar.gz",
+        "https://github.com/openxla/xla/archive/20a482597b7dd3067b26ca382b88084ee5a21cf7.tar.gz",
     ],
 )
 
@@ -47,6 +47,7 @@ python_init_repositories(
         "3.10": "@xla//:requirements_lock_3_10.txt",
         "3.11": "@xla//:requirements_lock_3_11.txt",
         "3.12": "@xla//:requirements_lock_3_12.txt",
+        "3.13": "@xla//:requirements_lock_3_13.txt",
     },
 )
 
@@ -82,6 +83,53 @@ load("@xla//:workspace0.bzl", "xla_workspace0")
 
 xla_workspace0()
 
+load(
+    "@xla//third_party/tsl/third_party/gpus/cuda/hermetic:cuda_json_init_repository.bzl",
+    "cuda_json_init_repository",
+)
+
+cuda_json_init_repository()
+
+load(
+    "@cuda_redist_json//:distributions.bzl",
+    "CUDA_REDISTRIBUTIONS",
+    "CUDNN_REDISTRIBUTIONS",
+)
+load(
+    "@xla//third_party/tsl/third_party/gpus/cuda/hermetic:cuda_redist_init_repositories.bzl",
+    "cuda_redist_init_repositories",
+    "cudnn_redist_init_repository",
+)
+
+cuda_redist_init_repositories(
+    cuda_redistributions = CUDA_REDISTRIBUTIONS,
+)
+
+cudnn_redist_init_repository(
+    cudnn_redistributions = CUDNN_REDISTRIBUTIONS,
+)
+
+load(
+    "@xla//third_party/tsl/third_party/gpus/cuda/hermetic:cuda_configure.bzl",
+    "cuda_configure",
+)
+
+cuda_configure(name = "local_config_cuda")
+
+load(
+    "@xla//third_party/tsl/third_party/nccl/hermetic:nccl_redist_init_repository.bzl",
+    "nccl_redist_init_repository",
+)
+
+nccl_redist_init_repository()
+
+load(
+    "@xla//third_party/tsl/third_party/nccl/hermetic:nccl_configure.bzl",
+    "nccl_configure",
+)
+
+nccl_configure(name = "local_config_nccl")
+
 load(
     "@bazel_toolchains//repositories:repositories.bzl",
     bazel_toolchains_repositories = "repositories",

docs/acc_jax.md

@@ -1,15 +1,16 @@
 # Accelerated JAX on Intel GPU
 
 ## Intel® Extension for OpenXLA* plug-in
-Intel® Extension for OpenXLA* includes PJRT plugin implementation, which seamlessly runs JAX models on Intel GPU. The PJRT API simplified the integration, which allowed the Intel GPU plugin to be developed separately and quickly integrated into JAX. Refer to [OpenXLA PJRT Plugin RFC](https://github.com/openxla/community/blob/main/rfcs/20230123-pjrt-plugin.md) for more details.
+Intel® Extension for OpenXLA includes PJRT plugin implementation, which seamlessly runs JAX models on Intel GPU. The PJRT API simplifies the integration, enabling the Intel GPU plugin to be developed separately and allowing smooth integration with JAX. This same PJRT implementation also enables initial Intel GPU support for TensorFlow and PyTorch models with XLA acceleration. Refer to [OpenXLA PJRT Plugin RFC](https://github.com/openxla/community/blob/main/rfcs/20230123-pjrt-plugin.md) for more details.
 
 ## Requirements
-Please check [README#requirements](../README.md#2-requirements) for the requirements of hardware and software.
+Please check the [Requirements section](../README.md#2-requirements) for the hardware and software requirements.
 
 ## Install
 The following table tracks intel-extension-for-openxla versions and compatible versions of `jax` and `jaxlib`. The compatibility between `jax` and `jaxlib` is maintained through JAX. This version restriction will be relaxed over time as the plugin API matures.
 |**intel-extension-for-openxla**|**jaxlib**|**jax**|
 |:-:|:-:|:-:|
+| 0.6.0 | 0.4.38 | 0.4.38 |
 | 0.5.0 | 0.4.30 | >= 0.4.30, <= 0.4.31|
 | 0.4.0 | 0.4.26 | >= 0.4.26, <= 0.4.27|
 | 0.3.0 | 0.4.24 | >= 0.4.24, <= 0.4.27|
@@ -35,7 +36,7 @@ python -c "import jax; print(jax.devices())"
 ```
 Reference result:
 ```
-[xpu(id=0), xpu(id=1)]
+[sycl(id=0), sycl(id=1)]
 ```
 
 ## Example - Run Stable Diffusion Inference
@@ -53,8 +54,8 @@ pip install transformers==4.47 diffusers==0.31.0 datasets==4.9.7 msgpack==1.1.0
 ```
 Source OneAPI env
 ```
-source /opt/intel/oneapi/compiler/2025.0/env/vars.sh
-source /opt/intel/oneapi/mkl/2025.0/env/vars.sh
+source /opt/intel/oneapi/compiler/2025.1/env/vars.sh
+source /opt/intel/oneapi/mkl/2025.1/env/vars.sh
 export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/intel/oneapi/umf/latest/lib
 ```
 **NOTE**: The path of OneAPI env script is based on the OneAPI installed path.
@@ -89,5 +90,5 @@ To submit questions, feature requests, and bug reports about the intel-extension
 2. If there is an error 'version GLIBCXX_3.4.30' not found, upgrade libstdc++ to the latest, for example for conda
 
     ```bash
-    conda install libstdcxx-ng==12.2.0 -c conda-forge
+    conda install libstdcxx-ng -c conda-forge -y
     ```

example/bert/requirements.txt

@@ -1,4 +1,4 @@
 datasets==2.20.0
 optax>=0.0.8
-transformers==4.48
+transformers>=4.48.0
 evaluate>=0.4.1

example/gemma/README.md

@@ -93,7 +93,7 @@ pip install -U tensorflow-text
 
 ### Examples
 
-#### Parameters of Single GPU finetune.
+#### Parameters for Single GPU finetune.
 ```
 python finetune.py \
 --model gemma_7b \
@@ -109,7 +109,7 @@ python finetune.py \
 --lora_rank 4 
 ```
 
-#### Parameters of Multiple GPUs data parallel distributed finetune
+#### Parameters for Multiple GPUs data parallel distributed finetune
 
 ```
 python finetune.py \

example/gptj/README.md

@@ -7,7 +7,7 @@ Script jax_gptj.py for [EleutherAI/gpt-j-6B](https://huggingface.co/EleutherAI/g
 Mark `intel-extension-for-openxla` folder as \<WORKSPACE\>, then
 ```bash
 cd <WORKSPACE>/example/gptj/
-pip install transformers==4.38 datasets==2.20.0
+pip install transformers==4.48 datasets==2.20.0
 pip install -r ../../test/requirements.txt
 ```
 

example/stable_diffusion/README.md

@@ -13,7 +13,7 @@ please got the [main page](https://github.com/intel/intel-extension-for-openxla/
 Mark `intel-extension-for-openxla` folder as \<WORKSPACE\>, then
 ```bash
 cd <WORKSPACE>/example/stable_diffusion/
-pip install transformers==4.38 diffusers==0.26.3 datasets==2.20.0 msgpack==1.0.7
+pip install transformers==4.48 diffusers==0.26.3 datasets==2.20.0 msgpack==1.0.7
 pip install -r ../../test/requirements.txt
 ```
 

example/t5/install_xpu.sh

@@ -10,7 +10,7 @@ pip uninstall tensorflow-metadata numba cudf -y
 pip uninstall tensorflow -y
 pip install tensorflow==2.18.0
 
-conda install libstdcxx-ng==12.2.0 -c conda-forge -y
+conda install libstdcxx-ng -c conda-forge -y
 
 pip uninstall mdit-py-plugins jupytext -y
 pip install t5