intel/intel-optimized-xla

Intel® Optimized OpenXLA*

Transformable numerical computing at scale combined with https://github.com/intel/intel-extension-for-openxla, which includes a PJRT plugin implementation to seamlessly runs https://github.com/google/jax models on Intel GPUs.

Images

The images below include https://github.com/google/jax and https://github.com/intel/intel-extension-for-openxla.

The images below additionally include Jupyter Notebook server:

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.6.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*:

The images below additionally include Jupyter Notebook server:

Build from Source

To build the images from source, clone the 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:

License

View the 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.