jax-ml / jax

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
http://jax.readthedocs.io/
Apache License 2.0
30.64k stars 2.82k forks source link

Intel (i)GPU support #22767

Open arunoruto opened 4 months ago

arunoruto commented 4 months ago

Currently supported platforms are CPUs, TPUs, and Nvidia GPUs, with AMD GPUs and the Metal API for MacOS being experimental. I was wondering if Intel Arc and iGPUs will be supported in the future. Looking at juliagpu, they support all of the platforms above (excluding TPUs) and utilize Intels OneAPI for computing.

MasterSkepticista commented 3 months ago

You can use JAX with Intel GPUs via intel-extension-for-openxla