diff --git a/.gitignore b/.gitignore
index f3e1d08f..d6f27a92 100644
--- a/.gitignore
+++ b/.gitignore
@@ -14,4 +14,5 @@ logs/
models-perf/
output/
site
+test-runner-summary-output.json
venv/
diff --git a/docs/roadmap.md b/docs/roadmap.md
index 8e22e7c8..018808b6 100644
--- a/docs/roadmap.md
+++ b/docs/roadmap.md
@@ -10,7 +10,7 @@
- Granite Rapids Support
- CLS Support
-- Intel Developer Cloud Support
+- Intel Tiber Developer Cloud Support
- AI Tools 2024.3/2025.0 Support
## Q4'24
diff --git a/docs/scripts/hook.py b/docs/scripts/hook.py
index 3b862bdf..2f0c96ec 100644
--- a/docs/scripts/hook.py
+++ b/docs/scripts/hook.py
@@ -34,6 +34,7 @@ def create_support_matrix():
compose_to_csv("pytorch", "serving")
compose_to_csv("tensorflow", None)
compose_to_csv("classical-ml", None)
+ compose_to_csv("jax", None)
# get_repo(models)
compose_to_csv("preset/data-analytics", "data_analytics")
diff --git a/docs/scripts/readmes.py b/docs/scripts/readmes.py
index 3e7d5e09..8eb2553b 100644
--- a/docs/scripts/readmes.py
+++ b/docs/scripts/readmes.py
@@ -17,6 +17,7 @@
readmes = [
"classical-ml/README.md",
+ "jax/README.md",
"preset/README.md",
"python/README.md",
"pytorch/README.md",
diff --git a/jax/.actions.json b/jax/.actions.json
new file mode 100644
index 00000000..36e21ad8
--- /dev/null
+++ b/jax/.actions.json
@@ -0,0 +1,5 @@
+{
+ "PACKAGE_OPTION": ["idp", "pip"],
+ "experimental": [true],
+ "runner_label": ["PVC"]
+}
diff --git a/jax/Dockerfile b/jax/Dockerfile
new file mode 100644
index 00000000..0d0118fa
--- /dev/null
+++ b/jax/Dockerfile
@@ -0,0 +1,104 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+#
+#
+# This file was assembled from multiple pieces, whose use is documented
+# throughout. Please refer to the TensorFlow dockerfiles documentation
+# for more information.
+
+ARG REGISTRY
+ARG REPO
+ARG GITHUB_RUN_NUMBER
+ARG BASE_IMAGE_NAME
+ARG BASE_IMAGE_TAG
+ARG PACKAGE_OPTION=pip
+ARG PYTHON_VERSION
+ARG PYTHON_BASE=${REGISTRY}/${REPO}:b-${GITHUB_RUN_NUMBER}-${BASE_IMAGE_NAME}-${BASE_IMAGE_TAG}-${PACKAGE_OPTION}-py${PYTHON_VERSION}-base
+ARG TORCHSERVE_BASE=${PYTHON_BASE}
+FROM ${PYTHON_BASE} AS xpu-base
+
+RUN apt-get update && \
+ apt-get install -y --no-install-recommends --fix-missing \
+ apt-utils \
+ build-essential \
+ clinfo \
+ git \
+ gnupg2 \
+ gpg-agent \
+ rsync \
+ unzip && \
+ apt-get clean && \
+ rm -rf /var/lib/apt/lists/*
+
+RUN wget -qO - https://repositories.intel.com/gpu/intel-graphics.key | \
+ gpg --dearmor --yes --output /usr/share/keyrings/intel-graphics.gpg
+RUN echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/gpu/ubuntu jammy unified" | \
+ tee /etc/apt/sources.list.d/intel-gpu-jammy.list
+
+ARG ICD_VER
+ARG LEVEL_ZERO_GPU_VER
+ARG LEVEL_ZERO_VER
+ARG LEVEL_ZERO_DEV_VER
+
+RUN apt-get update && \
+ apt-get install -y --no-install-recommends --fix-missing \
+ intel-opencl-icd=${ICD_VER} \
+ intel-level-zero-gpu=${LEVEL_ZERO_GPU_VER} \
+ libze1=${LEVEL_ZERO_VER} \
+ libze-dev=${LEVEL_ZERO_DEV_VER} && \
+ rm -rf /var/lib/apt/lists/*
+
+RUN no_proxy="" NO_PROXY="" wget --progress=dot:giga -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB \
+ | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
+ echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" \
+ | tee /etc/apt/sources.list.d/oneAPI.list
+
+ARG DPCPP_VER
+ARG MKL_VER
+ARG CCL_VER
+
+RUN apt-get update && \
+ apt-get install -y --no-install-recommends --fix-missing \
+ intel-oneapi-runtime-dpcpp-cpp=${DPCPP_VER} \
+ intel-oneapi-runtime-mkl=${MKL_VER} \
+ intel-oneapi-runtime-ccl=${CCL_VER} && \
+ rm -rf /var/lib/apt/lists/*
+
+RUN rm -rf /etc/apt/sources.list.d/intel-gpu-jammy.list /etc/apt/sources.list.d/oneAPI.list
+
+ENV OCL_ICD_VENDORS=/etc/OpenCL/vendors
+
+FROM xpu-base AS jax-base
+
+WORKDIR /
+COPY requirements.txt .
+
+RUN python -m pip install --no-cache-dir -r requirements.txt && \
+ rm -rf requirements.txt
+
+FROM jax-base AS jupyter
+
+WORKDIR /jupyter
+COPY jupyter-requirements.txt .
+
+RUN python -m pip install --no-cache-dir -r jupyter-requirements.txt && \
+ rm -rf jupyter-requirements.txt
+
+RUN mkdir -p /jupyter/ && chmod -R a+rwx /jupyter/
+RUN mkdir /.local && chmod a+rwx /.local
+
+EXPOSE 8888
+
+CMD ["bash", "-c", "source /etc/bash.bashrc && jupyter notebook --notebook-dir=/jupyter --port 8888 --ip 0.0.0.0 --no-browser --allow-root --ServerApp.token= --ServerApp.password= --ServerApp.allow_origin=* --ServerApp.base_url=$NB_PREFIX"]
diff --git a/jax/README.md b/jax/README.md
new file mode 100644
index 00000000..67ea81ed
--- /dev/null
+++ b/jax/README.md
@@ -0,0 +1,86 @@
+# Intel® Optimized OpenXLA\*
+
+Transformable numerical computing at scale combined with [Intel® Extension for OpenXLA\*], which includes a PJRT plugin implementation to seamlessly runs [JAX\*] models on Intel GPUs.
+
+## Images
+
+The images below include [JAX\*] and [Intel® Extension for OpenXLA\*].
+
+| Tag(s) | [JAX\*] | [Intel® Extension for OpenXLA\*] | [Flax] | Dockerfile |
+| -------------------------- | --------- | ----------------- | -------- | --------------- |
+| `0.4.0-pip-base`, `latest` | [v0.4.32] | [v0.4.0-jax] | [v0.9.0] | [v0.4.0] |
+
+The images below additionally include [Jupyter Notebook](https://jupyter.org/) server:
+
+| Tag(s) | [JAX\*] | [Intel® Extension for OpenXLA\*] | [Flax] | Dockerfile |
+| ------------------- | --------- | ----------------- | -------- | --------------- |
+| `0.4.0-pip-jupyter` | [v0.4.32] | [v0.4.0-jax] | [v0.9.0] | [v0.4.0] |
+
+### Run the Jupyter Container
+
+```bash
+docker run -it --rm \
+ -p 8888:8888 \
+ --net=host \
+ -v $PWD/workspace:/workspace \
+ -w /workspace \
+ intel/intel-optimized-xla:0.4.0-pip-jupyter
+```
+
+After running the command above, copy the URL (something like `http://127.0.0.1:$PORT/?token=***`) into your browser to access the notebook server.
+
+## Images with Intel® Distribution for Python*
+
+The images below include [Intel® Distribution for Python*]:
+
+| Tag(s) | [JAX\*] | [Intel® Extension for OpenXLA\*] | [Flax] | Dockerfile |
+| ---------------- | --------- | ----------------- | -------- | --------------- |
+| `0.4.0-idp-base` | [v0.4.32] | [v0.4.0-jax] | [v0.9.0] | [v0.4.0] |
+
+The images below additionally include [Jupyter Notebook](https://jupyter.org/) server:
+
+| Tag(s) | [JAX\*] | [Intel® Extension for OpenXLA\*] | [Flax] | Dockerfile |
+| ------------------- | --------- | ----------------- | -------- | --------------- |
+| `0.4.0-idp-jupyter` | [v0.4.32] | [v0.4.0-jax] | [v0.9.0] | [v0.4.0] |
+
+## Build from Source
+
+To build the images from source, clone the [AI Containers](https://github.com/intel/ai-containers) repository, follow the main `README.md` file to setup your environment, and run the following command:
+
+```bash
+cd jax
+docker compose build jax-base
+docker compose run -it jax-base
+```
+
+You can find the list of services below for each container in the group:
+
+| Service Name | Description |
+| ------------ | ----------------------------------------------- |
+| `jax-base` | Base image with [Intel® Extension for OpenXLA\*] |
+| `jupyter` | Adds Jupyter Notebook server |
+
+## License
+
+View the [License](https://github.com/intel/ai-containers/blob/main/LICENSE) for the [Intel® Distribution for Python].
+
+The images below also contain other software which may be under other licenses (such as Pytorch*, Jupyter*, Bash, etc. from the base).
+
+It is the image user's responsibility to ensure that any use of The images below comply with any relevant licenses for all software contained within.
+
+\* Other names and brands may be claimed as the property of others.
+
+<!--Below are links used in these document. They are not rendered: -->
+
+[Intel® Distribution for Python*]: https://www.intel.com/content/www/us/en/developer/tools/oneapi/distribution-for-python.html#gs.9bos9m
+[Intel® Extension for OpenXLA\*]: https://github.com/intel/intel-extension-for-openxla
+[JAX\*]: https://github.com/google/jax
+[Flax]: https://github.com/google/flax
+
+[v0.4.32]: https://github.com/google/jax/releases/tag/jax-v0.4.32
+
+[v0.4.0-jax]: https://github.com/intel/intel-extension-for-openxla/releases/tag/0.4.0
+
+[v0.9.0]: https://github.com/google/Flax/releases/tag/v0.9.0
+
+[v0.4.0]: https://github.com/intel/ai-containers/blob/v0.4.0/jax/Dockerfile
diff --git a/jax/docker-compose.yaml b/jax/docker-compose.yaml
new file mode 100644
index 00000000..e2c47d63
--- /dev/null
+++ b/jax/docker-compose.yaml
@@ -0,0 +1,87 @@
+# Copyright (c) 2024 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+include:
+ - path:
+ - ../python/docker-compose.yaml
+services:
+ jax-base:
+ build:
+ args:
+ http_proxy: ${http_proxy}
+ https_proxy: ${https_proxy}
+ no_proxy: ""
+ BASE_IMAGE_NAME: ${BASE_IMAGE_NAME:-ubuntu}
+ BASE_IMAGE_TAG: ${BASE_IMAGE_TAG:-22.04}
+ CCL_VER: ${CCL_VER:-2021.13.1-31}
+ DPCPP_VER: ${DPCPP_VER:-2024.2.1-1079}
+ GITHUB_RUN_NUMBER: ${GITHUB_RUN_NUMBER:-0}
+ ICD_VER: ${ICD_VER:-24.22.29735.27-914~22.04}
+ LEVEL_ZERO_DEV_VER: ${LEVEL_ZERO_DEV_VER:-1.17.6-914~22.04}
+ LEVEL_ZERO_GPU_VER: ${LEVEL_ZERO_GPU_VER:-1.3.29735.27-914~22.04}
+ LEVEL_ZERO_VER: ${LEVEL_ZERO_VER:-1.17.6-914~22.04}
+ MINIFORGE_VERSION: ${MINIFORGE_VERSION:-Linux-x86_64}
+ MKL_VER: ${MKL_VER:-2024.2.1-103}
+ NO_PROXY: ''
+ PACKAGE_OPTION: ${PACKAGE_OPTION:-pip}
+ PYTHON_VERSION: ${PYTHON_VERSION:-3.10}
+ REGISTRY: ${REGISTRY}
+ REPO: ${REPO}
+ context: .
+ labels:
+ dependency.python: ${PYTHON_VERSION:-3.10}
+ dependency.apt.build-essential: true
+ dependency.apt.clinfo: true
+ dependency.apt.git: true
+ dependency.apt.gnupg2: true
+ dependency.apt.gpg-agent: true
+ dependency.apt.intel-level-zero-gpu: ${LEVEL_ZERO_GPU_VER:-1.3.29735.27-914~22.04}
+ dependency.apt.intel-oneapi-runtime-ccl: ${CCL_VER:-2021.13.1-31}
+ dependency.apt.intel-oneapi-runtime-dpcpp-cpp: ${DPCPP_VER:-2024.2.1-1079}
+ dependency.apt.intel-oneapi-runtime-mkl: ${MKL_VER:-2024.2.1-103}
+ dependency.apt.intel-opencl-icd: ${ICD_VER:-23.43.27642.40-803~22.04}
+ dependency.apt.level-zero: ${LEVEL_ZERO_VER:-1.17.6-914~22.04}
+ dependency.apt.level-zero-dev: ${LEVEL_ZERO_DEV_VER:-1.17.6-914~22.04}
+ dependency.apt.rsync: true
+ dependency.apt.unzip: true
+ dependency.idp.pip: false
+ dependency.python.pip: requirements.txt
+ docs: jax
+ org.opencontainers.base.name: "intel/python:3.10-core"
+ org.opencontainers.image.name: "intel/intel-optimized-xla"
+ org.opencontainers.image.title: "Intel® Optimized XLA Base Image"
+ org.opencontainers.image.version: ${INTEL_XLA_VERSION:-v0.4.0}-${PACKAGE_OPTION:-pip}-base
+ target: jax-base
+ command: >
+ bash -c "python -c 'import jax; print(\"Jax Version:\", jax.__version__)'"
+ depends_on:
+ - ${PACKAGE_OPTION:-pip}
+ image: ${REGISTRY}/${REPO}:b-${GITHUB_RUN_NUMBER:-0}-${BASE_IMAGE_NAME:-ubuntu}-${BASE_IMAGE_TAG:-22.04}-${PACKAGE_OPTION:-pip}-py${PYTHON_VERSION:-3.10}-xla-${INTEL_XLA_VERSION:-v0.4.0}-base
+ pull_policy: always
+ jupyter:
+ build:
+ labels:
+ dependency.python.pip: jupyter-requirements.txt
+ org.opencontainers.base.name: "intel/intel-optimized-xla:${INTEL_XLA_VERSION:-v0.4.0}-base"
+ org.opencontainers.image.title: "Intel® Optimized XLA Jupyter Base Image"
+ org.opencontainers.image.version: ${INTEL_XLA_VERSION:-v0.4.0}-jupyter
+ target: jupyter
+ command: >
+ bash -c "python -m jupyter --version"
+ environment:
+ http_proxy: ${http_proxy}
+ https_proxy: ${https_proxy}
+ extends: jax-base
+ image: ${REGISTRY}/${REPO}:b-${GITHUB_RUN_NUMBER:-0}-${BASE_IMAGE_NAME:-ubuntu}-${BASE_IMAGE_TAG:-22.04}-${PACKAGE_OPTION:-pip}-py${PYTHON_VERSION:-3.10}-xla-${INTEL_XLA_VERSION:-v0.4.0}-jupyter
+ network_mode: host
diff --git a/jax/jupyter-requirements.txt b/jax/jupyter-requirements.txt
new file mode 100644
index 00000000..2cae0f91
--- /dev/null
+++ b/jax/jupyter-requirements.txt
@@ -0,0 +1,4 @@
+jupyterlab==4.2.4
+jupyterhub==5.1.0
+notebook==7.2.1
+jupyter-server-proxy>=4.1.2
diff --git a/jax/requirements.txt b/jax/requirements.txt
new file mode 100644
index 00000000..09d7cb7f
--- /dev/null
+++ b/jax/requirements.txt
@@ -0,0 +1,5 @@
+flax==0.8.2
+intel-extension-for-openxla==0.4.0
+jax==0.4.26
+jaxlib==0.4.26
+cython==3.0.11
diff --git a/jax/tests/example.py b/jax/tests/example.py
new file mode 100644
index 00000000..9227d066
--- /dev/null
+++ b/jax/tests/example.py
@@ -0,0 +1,37 @@
+# Copyright (c) 2024 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: skip-file
+
+import jax
+import jax.numpy as jnp
+
+print("jax.local_devices(): ", jax.local_devices())
+
+
+@jax.jit
+def lax_conv():
+ key = jax.random.PRNGKey(0)
+ lhs = jax.random.uniform(key, (2, 1, 9, 9), jnp.float32)
+ rhs = jax.random.uniform(key, (1, 1, 4, 4), jnp.float32)
+ side = jax.random.uniform(key, (1, 1, 1, 1), jnp.float32)
+ out = jax.lax.conv_with_general_padding(
+ lhs, rhs, (1, 1), ((0, 0), (0, 0)), (1, 1), (1, 1)
+ )
+ out = jax.nn.relu(out)
+ out = jnp.multiply(out, side)
+ return out
+
+
+print(lax_conv())
diff --git a/jax/tests/tests.yaml b/jax/tests/tests.yaml
new file mode 100644
index 00000000..419dbf3f
--- /dev/null
+++ b/jax/tests/tests.yaml
@@ -0,0 +1,29 @@
+# Copyright (c) 2024 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+---
+jax-import-${PACKAGE_OPTION:-pip}:
+ img: ${REGISTRY}/${REPO}:b-${GITHUB_RUN_NUMBER:-0}-${BASE_IMAGE_NAME:-ubuntu}-${BASE_IMAGE_TAG:-22.04}-${PACKAGE_OPTION:-pip}-py${PYTHON_VERSION:-3.10}-xla-${INTEL_XLA_VERSION:-v0.4.0}-base
+ cmd: python -c 'import jax; print("Jax Version:", jax.__version__); print(jax.devices())'
+ device: ["/dev/dri"]
+jax-import-jupyter-${PACKAGE_OPTION:-pip}:
+ img: ${REGISTRY}/${REPO}:b-${GITHUB_RUN_NUMBER:-0}-${BASE_IMAGE_NAME:-ubuntu}-${BASE_IMAGE_TAG:-22.04}-${PACKAGE_OPTION:-pip}-py${PYTHON_VERSION:-3.10}-xla-${INTEL_XLA_VERSION:-v0.4.0}-jupyter
+ cmd: sh -c "python -m jupyter --version"
+jax-xpu-example-${PACKAGE_OPTION:-pip}:
+ img: ${REGISTRY}/${REPO}:b-${GITHUB_RUN_NUMBER:-0}-${BASE_IMAGE_NAME:-ubuntu}-${BASE_IMAGE_TAG:-22.04}-${PACKAGE_OPTION:-pip}-py${PYTHON_VERSION:-3.10}-xla-${INTEL_XLA_VERSION:-v0.4.0}-base
+ cmd: python /tests/example.py
+ device: ["/dev/dri"]
+ volumes:
+ - src: $PWD/jax/tests
+ dst: /tests