Contoh PJRT

Contoh: Plugin JAX CUDA

1. Implementasi PJRT C API melalui wrapper (pjrt_c_api_gpu.h).
2. Siapkan titik entri untuk paket (setup.py).
3. Terapkan metode initialize() (__init__.py).
4. Dapat diuji dengan pengujian jax apa pun untuk CUDA.

Penerapan Framework

Beberapa referensi untuk menggunakan PJRT di sisi framework, untuk berinteraksi dengan perangkat PJRT:

• JAX
  • jax-ml/jax berinteraksi dengan PJRT API melalui xla_client API
• GoMLX
  • gomlx/gopjrt
  • gomlx/gomlx > backends/xla
• ZML
  • Wrapper PJRT API pjrt.zig
  • Memuat context.zig Plugin PJRT
  • Berinteraksi dengan Buffer PJRT buffer.zig
  • Menjalankan modul melalui module.zig PJRT

Implementasi Hardware

• Plugin integrasi lengkap (PJRT+MLIR+XLA):
  • Plugin CPU XLA
  • Plugin GPU XLA
  • Plugin Intel XLA
• Plugin integrasi ringan (PJRT+MLIR):
  • Plugin StableHLO Reference Interpreter (plugin C++ berbasis MLIR, akan ditautkan setelah devlabs)
  • Plugin Tenstorrent-XLA (plugin C berbasis MLIR)