PJRT Examples

Example: JAX CUDA plugin

1. PJRT C API implementation through wrapper (pjrt_c_api_gpu.h).
2. Set up the entry point for the package (setup.py).
3. Implement an initialize() method (__init__.py).
4. Can be tested with any jax tests for CUDA.

Frameworks Implementations

Some references for using PJRT on the framework side, to interface with PJRT devices:

• JAX
  • jax-ml/jax interacts with PJRT APIs via the xla_client APIs
• GoMLX
  • gomlx/gopjrt
  • gomlx/gomlx > backends/xla
• ZML
  • PJRT API wrapper pjrt.zig
  • Load PJRT Plugin context.zig
  • Interacting with PJRT Buffers buffer.zig
  • Execute a module via PJRT module.zig

Hardware Implementations

• Full integration plugins (PJRT+MLIR+XLA):
  • XLA CPU Plugin
  • XLA GPU Plugin
  • Intel XLA Plugin
• Light integration plugins (PJRT+MLIR):
  • StableHLO Reference Interpreter plugin (MLIR-based, C++ plugin, to be linked after devlabs)
  • Tenstorrent-XLA plugin (MLIR-based, C plugin)