Skip to content

JAX Integrationlink

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 pjrt-plugin/ folder within the iree-samples repository for now, though it will likely move elsewhere as it matures.