JAX integration

Note

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

Overview

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 integrations/pjrt directory.