PJRT - Uniform Device API

PJRT C API is the uniform Device API that we want to add to the ML ecosystem. The long term vision is that: (1) frameworks (TF, JAX, etc.) will call PJRT, which has device-specific implementations that are opaque to the frameworks; (2) each device focuses on implementing PJRT APIs as PJRT plugins, which can be opaque to the frameworks.

Communication channel

• Please file issues in the OpenXla/xla repo.
• Join the pjrt-announcement maillist.

Resources

• PJRT C API header
• PJRT C API changelog
• PJRT C++ API Overview
• PJRT integration guide
• PJRT design docs
• PJRT API ABI versioning and compatibility
• PJRT Plugin Mechanism design doc
• OpenXLA/IREE PJRT plugin implementation