rebase 0.4.26 (#327)

Co-authored-by: Sheng, Yang <[email protected]>
Co-authored-by: fengqing.lu <[email protected]>
Co-authored-by: Lu Teng <[email protected]>
Co-authored-by: yisonzhu <[email protected]>
Co-authored-by: lingzhi98 <[email protected]>

Author: 6 people

README.md

@@ -13,9 +13,9 @@ This guide introduces the overview of OpenXLA high level integration structure a
 
 ## 1. Overview
 
- <p align="center">
-	 <img src="openxla_for_intel_gpu.jpg" width="50%">
- </p>
+<p align="center">
+    <img src="openxla_for_intel_gpu.jpg" width="50%">
+</p>
 
 * [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.
 * TensorFlow and PyTorch support is on the way.
@@ -37,8 +37,8 @@ Verified Hardware Platforms:
 * Ubuntu 22.04, SUSE Linux Enterprise Server(SLES) 15 SP4
   * Intel® Data Center GPU Max Series
 * Intel® oneAPI Base Toolkit 2024.1
-* Jax/Jaxlib 0.4.25
-* Python 3.9-3.11
+* 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.**
@@ -81,8 +81,9 @@ source /opt/intel/oneapi/mkl/2024.0/env/vars.sh
 ### Install Jax and Jaxlib
 
 ```bash
-pip install jax==0.4.25 jaxlib==0.4.25
+pip install -r test/requirements.txt
 ```
+Check [test/requirements.txt](test/requirements.txt) for more details.
 
 ## 3. Install
 
@@ -100,7 +101,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.3.0-cp39-cp39-linux_x86_64.whl
+pip install intel_extension_for_openxla-0.4.0-cp39-cp39-linux_x86_64.whl
 ```
 
 **Aditional Build Option**:

WORKSPACE

@@ -14,10 +14,10 @@ http_archive(
     name = "xla",
     patch_args = ["-p1"],
     patches = ["//third_party:openxla.patch"],
-    sha256 = "8a59b9af7d0850059d7043f7043c780066d61538f3af536e8a10d3d717f35089",
-    strip_prefix = "xla-4ccfe33c71665ddcbca5b127fefe8baa3ed632d4",
+    sha256 = "fa6e7d17acc362b56c57c43224e6e3eca8569adae864e2fa191cc9d13edf4309",
+    strip_prefix = "xla-4e8e23f16bc925b6f27817de098a8e1e81296bb5",
     urls = [
-        "https://github.com/openxla/xla/archive/4ccfe33c71665ddcbca5b127fefe8baa3ed632d4.tar.gz",
+        "https://github.com/openxla/xla/archive/4e8e23f16bc925b6f27817de098a8e1e81296bb5.tar.gz",
     ],
 )
 

example/bert/README.md

@@ -1,5 +1,5 @@
 # Quick Start for fine-tunes BERT on SQuAD
-Fine-tunes BERT model on SQuAD task by [Question Answering examples](https://github.com/huggingface/transformers/tree/v4.37.0/examples/flax/question-answering).
+Fine-tunes BERT model on SQuAD task by [Question Answering examples](https://github.com/huggingface/transformers/tree/v4.38.0/examples/flax/question-answering).
 This expample is referred from [HuggingFace Transformers](https://github.com/huggingface/transformers). See [Backup](#Backup) for modification details.
 
 ## Requirements
@@ -10,9 +10,9 @@ Please got the [main page](https://github.com/intel/intel-extension-for-openxla/
 ### 2. Install dependency
 Mark `intel-extension-for-openxla` folder as \<WORKSPACE\>, then
 ```bash
-pip install jax==0.4.25 jaxlib==0.4.25 flax==0.8.2
 cd <WORKSPACE>/example/bert
 pip install -r requirements.txt
+pip install -r ../../test/requirements.txt
 ```
 
 ### 3. Download pre-trained model

example/bert/requirements.txt

@@ -1,7 +1,4 @@
-datasets>=1.8.0
-jax==0.4.25
-jaxlib==0.4.25
-flax>=0.8.2
+datasets==2.20.0
 optax>=0.0.8
 transformers==4.38
 evaluate>=0.4.1

example/fp8/README.md

@@ -11,8 +11,11 @@ export KAGGLE_KEY=xxxxxxxx
 
 ## Package dependency
 
+Mark `intel-extension-for-openxla` folder as \<WORKSPACE\>, then
 ```bash
-pip install keras keras_nlp kagglehub
+cd <WORKSPACE>/example/fp8/
+pip install keras-nlp==0.10.0 keras==3.3.2 kagglehub==0.2.5
+pip install -r ../../test/requirements.txt
 ```
 
 ## Dataset

example/gemma/README.md

@@ -13,8 +13,11 @@ export KAGGLE_KEY=xxxxxxxx
 
 ### Prerequisites
 
+Mark `intel-extension-for-openxla` folder as \<WORKSPACE\>, then
 ```bash
-pip install jax==0.4.25 jaxlib==0.4.25 keras-nlp==0.10.0 keras==3.3.2
+cd <WORKSPACE>/example/gemma/
+pip install keras-nlp==0.10.0 keras==3.3.2
+pip install -r ../../test/requirements.txt
 ```
 
 ### Options

example/gptj/README.md

@@ -4,8 +4,11 @@ Script jax_gptj.py for [EleutherAI/gpt-j-6B](https://huggingface.co/EleutherAI/g
 
 ## Prerequisites
 
+Mark `intel-extension-for-openxla` folder as \<WORKSPACE\>, then
 ```bash
-pip install jax==0.4.25 jaxlib==0.4.25 flax==0.8.2 transformers==4.38 datasets==2.12.0
+cd <WORKSPACE>/example/gptj/
+pip transformers==4.38 datasets==2.20.0
+pip install -r ../../test/requirements.txt
 ```
 
 ## Options

example/grok/README.md

@@ -6,16 +6,17 @@ Loading and running the Grok-1 open-weights model by [Grok-1](https://github.com
 
 Please got the [main page](https://github.com/intel/intel-extension-for-openxla/blob/main/README.md#build-and-install), and follow the instructions to build and install intel-extension-for-openxla.
 
-## 2. Install jax and flax
-```bash
-pip install jax==0.4.25 jaxlib==0.4.25
-```
-## 3. Install dependency
+## 2. Install dependency
+
+Mark `intel-extension-for-openxla` folder as \<WORKSPACE\>, then
 ```bash
+cd <WORKSPACE>/example/grok/
 git clone https://github.com/xai-org/grok-1.git
 pip install -r grok-1/requirements.txt
+pip install -r ../../test/requirements.txt
 ```
-## 4. Download the weights
+
+## 3. Download the weights
 
 Please follow [Downloading the weights](https://github.com/xai-org/grok-1#downloading-the-weights) to get the weights.
 
@@ -30,7 +31,7 @@ grok-1/
 ├── ...
 ```
 
-## 5. Copy files & Run
+## 4. Copy files & Run
 ```bash
 cp prompt.json inference.py grok-1
 python grok-1/inference.py

example/resnet50/README.md

@@ -8,19 +8,20 @@ Trains ResNet50 model (He et al., 2016) for the ImageNet classification task (Ru
 
 Please got the [main page](https://github.com/intel/intel-extension-for-openxla/blob/main/README.md#build-and-install), and follow the instructions to build and install intel-extension-for-openxla.
 
-### 2. Install jax and flax
-```bash
-pip install jax==0.4.25 jaxlib==0.4.25
-```
-### 3. Install dependency
+### 2. Install dependency
+
+Mark `intel-extension-for-openxla` folder as \<WORKSPACE\>, then
 ```bash
+cd <WORKSPACE>/example/resnet50/
 git clone --branch=main https://github.com/google/flax
 cd flax
 git checkout ba9e24a7b697e6407465cb4b05a24a4cea152248
-cd examples/imagenet
-pip install -r requirements.txt
+pip install -r examples/imagenet/requirements.txt
+cd ..
+pip install -r ../../test/requirements.txt
 ```
-### 4. Download dataset
+
+### 3. Download dataset
 
 Please follow [Preparing the dataset](https://github.com/google/flax/tree/main/examples/imagenet#preparing-the-dataset) to get imagenet dataset.
 

example/sdxl/README.md

@@ -14,15 +14,15 @@ More details could be found in [Stability-AI/generative-models](https://github.c
 
 please got the [main page](https://github.com/intel/intel-extension-for-openxla/blob/main/README.md#build-and-install), and follow the instructions to build and install intel-extension-for-openxla.
 
-### 2. Install jax
-```bash
-pip install jax==0.4.25 jaxlib==0.4.25 flax==0.8.2
-```
-### 3. Install huggingface transformers
+### 2. Install packages
 
+Mark `intel-extension-for-openxla` folder as \<WORKSPACE\>, then
 ```bash
-pip install transformers==4.38 diffusers==0.26.3 datasets==2.12.0 msgpack==1.0.7
+cd <WORKSPACE>/example/sdxl/
+pip install transformers==4.38 diffusers==0.26.3 datasets==2.20.0 msgpack==1.0.7
+pip install -r ../../test/requirements.txt
 ```
+
 ## Run
 
 ### 1. Environmental Variables