JAX integration

!!! note IREE's JAX support is under active development. This page is still under construction.

IREE offers two ways to interface with JAX programs:

An API for extracting and compiling full models ahead of time (AOT) for execution apart from JAX. This API is being developed in the iree-org/iree-jax repository.
A PJRT plugin that adapts IREE as a native JAX backend for online / just in time (JIT) use. This plugin is being developed in the openxla/openxla-pjrt-plugin repository.