intel/intel-optimized-xla

Intel® 优化的 OpenXLA*

提供可扩展的数值计算能力，结合https://github.com/intel/intel-extension-for-openxla，其中包含PJRT插件实现，可在Intel GPU上无缝运行https://github.com/google/jax模型。

镜像概述和主要用途

本镜像系列包含JAX*和Intel® OpenXLA扩展，提供基础运行环境、集成Jupyter Notebook的交互环境以及包含Intel® Python发行版的优化环境，适用于基于JAX框架的机器学习模型开发、训练和部署，特别优化了在Intel GPU上的运行性能。

核心功能和特性

• 集成JAX*深度学习框架，支持高性能数值计算
• 包含Intel® OpenXLA扩展，通过PJRT插件实现Intel GPU加速
• 提供Flax库支持，便于构建神经网络
• 部分版本集成Jupyter Notebook，支持交互式开发
• 提供基于Intel® Python发行版的优化版本，提升计算性能

镜像标签及组件版本

基础镜像（含JAX和Intel® OpenXLA扩展）

含Jupyter Notebook的镜像

含Intel® Python发行版的镜像

含Intel® Python发行版和Jupyter Notebook的镜像

使用方法

运行Jupyter容器

对于包含Jupyter Notebook的镜像，可通过以下命令启动：

bash
docker run -it --rm \
    -p 8888:8888 \
    --net=host \
    -v $PWD/workspace:/workspace \
    -w /workspace \
    docker.xuanyuan.run/intel/intel-optimized-xla:0.6.0-pip-jupyter

运行上述命令后，将URL（类似http://127.0.0.1:$PORT/?token=）复制到浏览器中以访问notebook服务器。

从源码构建

要从源码构建镜像，请克隆https://github.com/intel/ai-containers仓库，按照主`README.md`文件设置环境，然后运行以下命令：

bash
cd jax
docker compose build jax-base
docker compose run -it jax-base

容器组中各服务说明：

许可证

查看[Intel® Python发行版]的https://github.com/intel/ai-containers/blob/main/LICENSE。

本镜像还包含其他可能受其他许可证约束的软件（如基础镜像中的Pytorch*、Jupyter*、Bash等）。

镜像用户有责任确保对本镜像的任何使用均符合其中包含的所有软件的相关许可证要求。

* 其他名称和品牌可能是其各自所有者的财产。