דוגמאות של PJRT

דוגמה: הפלאגין JAX CUDA

1. הטמעת PJRT C API דרך wrapper (pjrt_c_api_gpu.h).
2. מגדירים את נקודת הכניסה לחבילה (setup.py).
3. מטמיעים את השיטה initialize()‎ (__init__.py).
4. אפשר לבדוק אותו באמצעות כל בדיקות jax ל-CUDA.

הטמעות של Frameworks

מקורות מידע נוספים לשימוש ב-PJRT בצד המסגרת, כדי ליצור ממשק עם מכשירי PJRT:

• JAX
  • jax-ml/jax מקיים אינטראקציה עם ממשקי ה-API של PJRT דרך ממשקי ה-API של xla_client
• GoMLX
  • gomlx/gopjrt
  • gomlx/gomlx > backends/xla
• ZML
  • מעטפת PJRT API‏ pjrt.zig
  • טעינת הפלאגין PJRT‏ context.zig
  • אינטראקציה עם מאגרי PJRT‏ buffer.zig
  • הפעלת מודול באמצעות PJRT module.zig

הטמעות בחומרה

• פלאגינים לשילוב מלא (PJRT+‏MLIR+‏XLA):
  • Plugin של XLA CPU
  • Plugin של XLA GPU
  • הפלאגין Intel XLA
• פלאגינים לשילוב קל (PJRT+MLIR):
  • הפלאגין של StableHLO Reference Interpreter (פלאגין מבוסס-MLIR ב-C++, שצריך לקשר אחרי devlabs)
  • הפלאגין Tenstorrent-XLA (פלאגין C שמבוסס על MLIR)