Skip to content

JAX integrationlink


IREE's JAX support is under 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 integrations/pjrt directory.