Support for jax 0.8.0

[llandsmeer]

Hi,

The table in the readme lists jax 0.4.38 as the latest supported version.

Are there any plans for officially supporting the latest jax version 0.8.0?

Thanks,
Lennart

[vsanghavi]

Hello! We’re making steady progress on closing the gap between the JAX version compatible with Intel GPUs and the versions released upstream. For the upcoming Intel Extension for OpenXLA 0.7.0 release, JAX v0.5.0 will be the supported version.

Starting next year, we plan to introduce a JAX plugin for Intel devices that allows users to run JAX directly from the upstream project, without requiring the Intel Extension for OpenXLA. This will enable us to support versions that are close to—or fully aligned with—the latest upstream JAX releases.

To better understand your needs and guide our roadmap, could you share which Intel GPU you’re currently using or intend to use in the future?