Rebase with jax-v0.4.20 (#119)

Author: Solaryee

.bazelrc

@@ -24,6 +24,7 @@ build:gpu --define=using_sycl=true
 build:gpu --define=device=gpu
 build:gpu --config=onednn_v3
 build:gpu --repo_env TF_NEED_SYCL=1
+build:gpu --define=tensorflow_mkldnn_contraction_kernel=0
 
 # This config build with oneDNN V3 API, which is enabled by default
 build:onednn_v3 --copt=-DITEX_ONEDNN_3_0 --define=onednn_version=3

README.md

@@ -37,7 +37,7 @@ Verified Hardware Platforms:
 * Ubuntu 22.04, Red Hat 8.6/8.8/9.2 (64-bit), SUSE Linux Enterprise Server(SLES) 15 SP4
   * Intel® Data Center GPU Max Series
 * Intel® oneAPI Base Toolkit 2023.2
-* Jax/Jaxlib 0.4.13
+* Jax/Jaxlib 0.4.20
 * Python 3.9-3.11
 * pip 19.0 or later (requires manylinux2014 support)
 
@@ -69,7 +69,7 @@ source /opt/intel/oneapi/tbb/2021.9.0/env/vars.sh
 ### Install Jax and Jaxlib
 
 ```bash
-pip install jax==0.4.13 jaxlib==0.4.13
+pip install jax==0.4.20 jaxlib==0.4.20
 ```
 
 ## 3. Install

WORKSPACE

@@ -14,10 +14,10 @@ http_archive(
     name = "xla",
     patch_args = ["-p1"],
     patches = ["//third_party:openxla.patch"],
-    sha256 = "4ec16aff3862c5a243db956ce558d7a62eb79f5e20747b0e80802a3b0d12e419",
-    strip_prefix = "xla-12de6ec958419b57be248d0acd2d9f757e71748c",
+    sha256 = "15235814a637f2199d2adae05c09dc8f6655737f06713108f970749bc5e87b46",
+    strip_prefix = "xla-ca31652cdbeb6ea187589dea546ff8019274f8b2",
     urls = [
-        "https://github.com/openxla/xla/archive/12de6ec958419b57be248d0acd2d9f757e71748c.tar.gz",
+        "https://github.com/openxla/xla/archive/ca31652cdbeb6ea187589dea546ff8019274f8b2.tar.gz",
     ],
 )
 
@@ -62,4 +62,4 @@ bazel_toolchains_repositories()
 
 load("//xla:workspace.bzl", "workspace")
 
-workspace()
+workspace()

example/bert/README.md

@@ -10,7 +10,7 @@ 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.13 jaxlib==0.4.13 flax==0.7.0
+pip install jax==0.4.20 jaxlib==0.4.20 flax==0.7.0
 cd <WORKSPACE>/example/bert
 pip install -r requirements.txt
 ```

example/gptj/README.md

@@ -5,7 +5,7 @@ Script jax_gptj.py for [EleutherAI/gpt-j-6B](https://huggingface.co/EleutherAI/g
 ## Prerequisites
 
 ```bash
-pip install jax==0.4.13 jaxlib==0.4.13 flax==0.7.0 transformers==4.27.4 diffusers==0.16.1 datasets==2.12.0
+pip install jax==0.4.20 jaxlib==0.4.20 flax==0.7.0 transformers==4.32 diffusers==0.16.1 datasets==2.12.0
 ```
 
 ## Options

example/resnet50/README.md

@@ -10,7 +10,7 @@ Please got the [main page](https://github.com/intel/intel-extension-for-openxla/
 
 ### 2. Install jax and flax
 ```bash
-pip install jax==0.4.13 jaxlib==0.4.13
+pip install jax==0.4.20 jaxlib==0.4.20
 ```
 ### 3. Install dependency
 ```bash

example/stable_diffusion/README.md

@@ -10,12 +10,12 @@ please got the [main page](https://github.com/intel/intel-extension-for-openxla/
 
 ### 2. Install jax
 ```bash
-pip install jax==0.4.13 jaxlib==0.4.13 flax==0.7.0
+pip install jax==0.4.20 jaxlib==0.4.20 flax==0.7.0
 ```
 ### 3. Install huggingface transformers
 
 ```bash
-pip install transformers==4.27.4 diffusers==0.16.1 datasets==2.12.0 msgpack==1.0.7
+pip install transformers==4.32 diffusers==0.16.1 datasets==2.12.0 msgpack==1.0.7
 ```
 ## Run
 
@@ -24,6 +24,7 @@ pip install transformers==4.27.4 diffusers==0.16.1 datasets==2.12.0 msgpack==1.0
 | **ENV** | **Description** | **PVC Platform** | **ATSM/DG2 Platform** | 
 | :---: | :---: | :---: |:---: |
 | ZE_AFFINITY_MASK | Run this model on single GPU tile |export ZE_AFFINITY_MASK=0.0 | export ZE_AFFINITY_MASK=0.0 | 
+| XLA_FLAGS | Cutomerlize xla debug options | export XLA_FLAGS="--xla_gpu_force_conv_nhwc" | export XLA_FLAGS="--xla_gpu_force_conv_nhwc" |
 
 ### 2. Inference Command
 

example/t5/README.md

@@ -37,7 +37,7 @@ For benchmarking, you could skip this step because our model script will downloa
 git clone https://github.com/google-research/t5x.git
 bash install_xpu.sh
 pip install --upgrade intel-extension-for-openxla
-pip install jax==0.4.13 jaxlib==0.4.13
+pip install jax==0.4.20 jaxlib==0.4.20
 ```
 ## Inference
 

example/t5/install_xpu.sh

@@ -13,7 +13,7 @@ pip uninstall tensorflow -y
 pip install tensorflow==2.12.0
 
 conda install libstdcxx-ng==12.2.0 -c conda-forge -y
-pip install jax==0.4.13 jaxlib==0.4.13
+pip install jax==0.4.20 jaxlib==0.4.20
 
 pip uninstall mdit-py-plugins jupytext -y
 pip install t5

example/t5/patch/version_time_dlpath.patch

@@ -8,8 +8,8 @@ index 37238ba..19bec2f 100644
 
 -_jax_version = '0.4.11'
 -_jaxlib_version = '0.4.11'
-+_jax_version = '0.4.13'
-+_jaxlib_version = '0.4.13'
++_jax_version = '0.4.20'
++_jaxlib_version = '0.4.20'
 
  setuptools.setup(
      name='t5x',