Release Notes

This is the first release of the oneAPI PJRT and plugin wheels for JAX on PyPI, enabling JAX applications to run on Intel® GPUs through XLA’s SYCL platform support. This release is based on JAX 0.11.2.

This is an experimental (Alpha) release. Not all JAX functionality is supported yet.

Features:

• Core JAX functionality: jax.jit, jax.grad/autodiff, jax.vmap, jax.numpy, jax.nn, jax.random (Threefry kernel), FFT, convolutions, sort & scan.
• GEMM in f16, bf16, f32, f64, c64, c128.
• LAPACK support for Cholesky, LU, solve / inverse / determinant, QR (incl. column pivoting), Eigenvalue decomposition (eigh), SVD, Eig with hybrid host/GPU execution, triangular solve.
• Single-host multi-GPU: pmap, jit with sharding, and collectives.
• Validated to work with MaxText, MaxDiffusion, Optax and Orbax,.

Compatibility

Hardware

• Validated on Intel® Arc™ Pro B50/ B60/B65/B70 GPUs and
• Intel® Data Center GPU Max Series.

Prerequisites

Intel GPU Driver. The intel GPU compute runtime should be installed on the system before installing the wheels. The pip wheels do not include them. Use one of the following options.

• Option 1: Container
  Use the Intel® Graphics Compute Runtime for oneAPI Level Zero and OpenCL™ Driver Container, (with tag 26.31.39395.13-runtime-<os version>) that contains the necessary driver components. Follow instructions on running the container.
• Option 2: Native Installation
  Install ONLY the intel-gpu-compute runtime from the Intel(R) OMIX repository. (OMIX 0.4). Follow the instructions for setting up the package repository and install the compute runtime with this command.

sudo apt install -y intel-gpu-compute

Installation

python -m venv jax-oneapi && source jax-oneapi/bin/activate
pip install "jax[oneapi]==0.11.2"
python -c "import jax; print(jax.devices())"
# Will print [OneapiDevice(id=0)]

jax[oneapi] installs jax, jaxlib, both oneAPI plugin wheels, and required dependencies.

Limitations

• The following features in JAX are currently not supported.
  • Scale-out – distribution across multiple GPUs across multiple hosts.
  • Advanced JAX features such as Pallas
  • Mosaic-GPU and other CUDA only features
  • Sparse kernels
  • RNN Kernels
• Known issues
  • Using JAX with other oneAPI based frameworks such as PyTorch can cause installation issues due to conflicts in the dependencies, in that case use each framework in a separate container or Python environment.
  • On the OneAPI backend, the missing-OneAPI-platform gap in XLA's serialization paths causes both executable serialization (UNIMPLEMENTED: Unsupported platform ID for XlaExecutableAbiVersion) and PJRT topology deserialization (INVALID_ARGUMENT: The platform is not a GPU platform. platform_name: 'oneapi') to fail.
  • Using jax.profiler.trace fails with “device-activity plane named /device:GPU:N” error.
  • DLPack has no kDLOneAPIHost device type, so OneAPI pinned-host memory round-trips incorrectly report memory_kind as 'device'; pinned-host DLPack round-trip is therefore unsupported.
  • On the OneAPI GPU backend, reverse-mode grad of odeint crashes during codegen with LLVM Fatal Error while compiling target binary for module: jit__odeint_wrapper failing test_pend_grads.
  • LLVM Fatal Error while compiling target binary for module: jit_convert_element_type due to unsupported SPIR-V extension: SPV_EXT_long_vector.
  • jnp.convolve / jnp.correlate with an unsigned-integer preferred_element_type fail to execute on the OneAPI GPU backend, raising INVALID_ARGUMENT: Unsupported element type: u16 and u32 when the result is materialized.
  • A subset of bfloat16 ops is miscomputed only when the buffer is a single element; Affects a few tests (testFull6, testOp61_exp_bfloat16, testReducerWhere62).
  • bfloat16 lax.sort mis-orders elements (e.g. testSortAgainstNumpy1, testSortNumKeys3) because vISA lacks BF16 immediate support, silently promoting operands and breaking bit-exact compares.
  • Due to lack of scale-out support, any multi-process (multi-host) collective spanning more than one process fails at runtime with JaxRuntimeError: INTERNAL: For non-local GPU cliques ... clique id callback must be passed via execution params, breaking the tests/multiprocess/* suite (psum/all_gather/all_to_all, pjit, array, callbacks, pmap).
  • The signed int4 binary ops (e.g. lax.min) returns wrong results; fails test_int4_binary_ops6.
  • On OneAPI GPU, the TF32_TF32_F32 and TF32_TF32_F32_X3 dot algorithm presets are not yet supported and will raise an error when used.
  • On the OneAPI backend, jax.export cannot serialize computations with OneAPI solver/linalg custom calls (targets lack export back-compat guarantees), failing shape_poly_test and export_harnesses_multi_platform_test; runtime execution of these ops is unaffected.
  • On the OneAPI GPU backend, jax.scipy.spatial.transform.Rotation.from_matrix(...).as_rotvec() crashes with a segmentation fault during compilation, failing tests/scipy_spatial_test.py::LaxBackedScipySpatialTransformTests.testRotationFromMatrix1
  • With shape-polymorphic (symbolic) dimensions, SVD fails at lowering time with jax._src.export.shape_poly.InconclusiveDimensionOperation on OneAPI, breaking the test_harness_svdpoly cases in tests/shape_poly_test.py; concrete-shape SVD is unaffected.
  • On OneAPI, the jitted NVFP4 scaled_dot_general_wrapper output doesn't match the explicitly dequantized float32 dot_general reference, failing the test_dot_general_nvfp4* cases in tests/scaled_matmul_stablehlo_test.py.