diff --git a/README.md b/README.md
index 7fff09f..c2d19e7 100644
--- a/README.md
+++ b/README.md
@@ -14,7 +14,7 @@ This guide introduces the overview of OpenXLA high level integration structure a
## 1. Overview
- 
+ 
* [JAX](https://jax.readthedocs.io/en/latest/) provides a familiar NumPy-style API, includes composable function transformations for compilation, batching, automatic differentiation, and parallelization, and the same code executes on multiple backends.
@@ -26,9 +26,9 @@ This guide introduces the overview of OpenXLA high level integration structure a
Verified Hardware Platforms:
-* Intel® Data Center GPU Max Series, Driver Version: [803](https://dgpu-docs.intel.com/releases/LTS_803.29_20240131.html)
+* Intel® Data Center GPU Max Series, Driver Version: [803](https://dgpu-docs.intel.com/releases/LTS_803.63_20240617.html)
-* Intel® Data Center GPU Flex Series 170, Driver Version: [803](https://dgpu-docs.intel.com/releases/LTS_803.29_20240131.html)
+* Intel® Data Center GPU Flex Series 170, Driver Version: [803](https://dgpu-docs.intel.com/releases/LTS_803.63_20240617.html)
### Software Requirements
@@ -36,19 +36,19 @@ Verified Hardware Platforms:
* 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 2024.1
+* [Intel® oneAPI Base Toolkit 2024.2](https://www.intel.com/content/www/us/en/developer/articles/release-notes/intel-oneapi-toolkit-release-notes.html)
* Jax/Jaxlib 0.4.26
* Python 3.9-3.12
* 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 is restricted when works 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/installation-guides/index.html#intel-data-center-gpu-flex-series) for latest driver installation. If install the verified Intel® Data Center GPU Max Series/Intel® Data Center GPU Flex Series [803](https://dgpu-docs.intel.com/releases/LTS_803.29_20240131.html), please append the specific version after components, such as `sudo apt-get install intel-opencl-icd==23.43.27642.38-803~22.04`|
-|Ubuntu 22.04, SLES 15 SP4|Intel® Data Center GPU Max Series| Refer to the [Installation Guides](https://dgpu-docs.intel.com/installation-guides/index.html#intel-data-center-gpu-max-series) for latest driver installation. If install the verified Intel® Data Center GPU Max Series/Intel® Data Center GPU Flex Series [803](https://dgpu-docs.intel.com/releases/LTS_803.29_20240131.html), please append the specific version after components, such as `sudo apt-get install intel-opencl-icd==23.43.27642.38-803~22.04`|
+|Ubuntu 22.04 |Intel® Data Center GPU Flex Series| Refer to the [Installation Guides](https://dgpu-docs.intel.com/installation-guides/index.html#intel-data-center-gpu-flex-series) for latest driver installation. If install the verified Intel® Data Center GPU Max Series/Intel® Data Center GPU Flex Series [803](https://dgpu-docs.intel.com/releases/LTS_803.63_20240617.html), please append the specific version after components, such as `sudo apt-get install intel-opencl-icd==23.43.27642.52-803~22.04`|
+|Ubuntu 22.04, SLES 15 SP4|Intel® Data Center GPU Max Series| Refer to the [Installation Guides](https://dgpu-docs.intel.com/installation-guides/index.html#intel-data-center-gpu-max-series) for latest driver installation. If install the verified Intel® Data Center GPU Max Series/Intel® Data Center GPU Flex Series [803](https://dgpu-docs.intel.com/releases/LTS_803.63_20240617.html), please append the specific version after components, such as `sudo apt-get install intel-opencl-icd==23.43.27642.52-803~22.04`|
### Install oneAPI Base Toolkit Packages
@@ -58,34 +58,34 @@ Need to install components of Intel® oneAPI Base Toolkit:
* Intel® oneAPI Math Kernel Library (oneMKL)
```bash
-$ wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/fdc7a2bc-b7a8-47eb-8876-de6201297144/l_BaseKit_p_2024.1.0.596.sh
+$ wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/e6ff8e9c-ee28-47fb-abd7-5c524c983e1c/l_BaseKit_p_2024.2.1.100_offline.sh
# 2 components are necessary: DPC++/C++ Compiler and oneMKL
-sudo sh l_BaseKit_p_2024.1.0.596.sh
+sudo sh l_BaseKit_p_2024.2.1.100_offline.sh
# Source OneAPI env
-source /opt/intel/oneapi/compiler/2024.1/env/vars.sh
-source /opt/intel/oneapi/mkl/2024.1/env/vars.sh
+source /opt/intel/oneapi/compiler/2024.2/env/vars.sh
+source /opt/intel/oneapi/mkl/2024.2/env/vars.sh
```
-**Backup**: Recommend to rollback to **Toolkit 2024.0** if need collective feature. See [Release Notes](https://github.com/intel-innersource/frameworks.ai.intel-extension-for-openxla.intel-extension-for-openxla/releases) for more details.
+**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/20f4e6a1-6b0b-4752-b8c1-e5eacba10e01/l_BaseKit_p_2024.0.0.49564.sh
+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.0.0.49564.sh
+sudo sh l_BaseKit_p_2024.1.0.596.sh
# Source OneAPI env
-source /opt/intel/oneapi/compiler/2024.0/env/vars.sh
-source /opt/intel/oneapi/mkl/2024.0/env/vars.sh
+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
-pip install -r test/requirements.txt
+pip install -r https://raw.githubusercontent.com/intel/intel-extension-for-openxla/main/test/requirements.txt
```
-Check [test/requirements.txt](test/requirements.txt) for more details.
+Please refer to [test/requirements.txt](test/requirements.txt) for the version dependency of `jax`, `jaxlib` and `flax`.
-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.
+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.4.0 | 0.4.26 | >= 0.4.26, <= 0.4.27|
diff --git a/docs/acc_jax.md b/docs/acc_jax.md
index 10a5f9f..64cd6e4 100644
--- a/docs/acc_jax.md
+++ b/docs/acc_jax.md
@@ -1,31 +1,13 @@
# 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 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.
## Requirements
-
-### Hardware Requirements
-
-Verified Hardware Platforms:
-
-* Intel® Data Center GPU Max Series, Driver Version: [803](https://dgpu-docs.intel.com/releases/LTS_803.63_20240617.html) ([Installation Guides](https://dgpu-docs.intel.com/installation-guides/index.html#intel-data-center-gpu-max-series))
-
-* Intel® Data Center GPU Flex Series 170, Driver Version: [803](https://dgpu-docs.intel.com/releases/LTS_803.63_20240617.html) ([Installation Guides](https://dgpu-docs.intel.com/installation-guides/index.html#intel-data-center-gpu-flex-series))
-
-### 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 2024.1 ([Installation Guides](https://github.com/intel/intel-extension-for-openxla/?tab=readme-ov-file#install-oneapi-base-toolkit-packages))
-* Jax/Jaxlib 0.4.26
-* Python 3.9-3.12
-* pip 19.0 or later (requires manylinux2014 support)
+Please check [README#requirements](../README.md#2-requirements) for the requirements of hardware and software.
## 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.
+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.4.0 | 0.4.26 | >= 0.4.26, <= 0.4.27|
@@ -39,9 +21,12 @@ The following table tracks intel-extension-for-openxla versions and compatible v
conda create -n jax-ioex python=3.10
conda activate jax-ioex
pip install -U pip
-pip install jax==0.4.26 jaxlib==0.4.26
pip install intel-extension-for-openxla
+
+# Install jax, jaxlib and flax dependency
+pip install -r https://raw.githubusercontent.com/intel/intel-extension-for-openxla/main/test/requirements.txt
```
+Please refer to [requirements.txt](../test/requirements.txt) for the version dependency of `jax`, `jaxlib` and `flax`.
## Verify
```
@@ -61,12 +46,8 @@ bash Miniforge3-$(uname)-$(uname -m).sh
```
### Setup environment
+Please follow [Install](#install) to prepare the basic environment first.
```
-conda create -n stable-diffusion python=3.10
-conda activate stable-diffusion
-pip install -U pip
-pip install jax==0.4.26 jaxlib==0.4.26 flax==0.8.2
-pip install intel-extension-for-openxla
pip install transformers==4.38 diffusers==0.26.3 datasets==2.12.0 msgpack==1.0.7
```
Source OneAPI env
@@ -76,9 +57,7 @@ source /opt/intel/oneapi/mkl/2024.2/env/vars.sh
```
**NOTE**: The path of OneAPI env script is based on the OneAPI installed path.
-
### Run Demo (Stable Diffusion Inference)
-
Go to [example/stable_diffusion](../example/stable_diffusion/README.md) for detail about this demo.
| **Command** | **Model** | **Output Image Resolution** |
diff --git a/example/gptj/README.md b/example/gptj/README.md
index 0c1c9eb..942bdcf 100644
--- a/example/gptj/README.md
+++ b/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 \, then
```bash
cd /example/gptj/
-pip transformers==4.38 datasets==2.20.0
+pip install transformers==4.38 datasets==2.20.0
pip install -r ../../test/requirements.txt
```
diff --git a/openxla_for_intel_gpu.jpg b/openxla_for_intel_gpu.jpg
deleted file mode 100644
index 0074978..0000000
Binary files a/openxla_for_intel_gpu.jpg and /dev/null differ
diff --git a/openxla_for_intel_gpu.png b/openxla_for_intel_gpu.png
new file mode 100644
index 0000000..01149b7
Binary files /dev/null and b/openxla_for_intel_gpu.png differ
diff --git a/xla/tools/pip_package/xla_setup.py b/xla/tools/pip_package/xla_setup.py
index 9ff252f..29cfbea 100644
--- a/xla/tools/pip_package/xla_setup.py
+++ b/xla/tools/pip_package/xla_setup.py
@@ -149,6 +149,7 @@ def finalize_options(self):
'Programming Language :: Python :: 3.9',
'Programming Language :: Python :: 3.10',
'Programming Language :: Python :: 3.11',
+ 'Programming Language :: Python :: 3.12',
'Topic :: Scientific/Engineering',
'Topic :: Scientific/Engineering :: Mathematics',
'Topic :: Scientific/Engineering :: Artificial Intelligence',