Project: /_project.yaml Book: /_book.yaml

<link rel="stylesheet" href="/site-assets/css/style.css">

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 channels

• Issues and feature requests can be filed in the OpenXLA/xla repo.
• Questions regarding PJRT can be asked on the OpenXLA Discord.

Resources

• PJRT C API header
• PJRT C API changelog
• PJRT integration guide
• PJRT design docs
• PJRT API ABI versioning and compatibility
• PJRT Plugin Mechanism design doc
• OpenXLA 2024 Fall DevLab PJRT plugin tutorial slides (recording)
• OpenXLA/IREE PJRT plugin implementation