مثال های PJRT

مثال: افزونه JAX CUDA

1. اجرای PJRT C API از طریق wrapper ( pjrt_c_api_gpu.h ).
2. نقطه ورود بسته را تنظیم کنید ( setup.py ).
3. یک متد ()initialize ( __init__.py ) را پیاده سازی کنید.
4. می توان با هر تست جکس برای CUDA تست کرد.

پیاده سازی چارچوب ها

برخی از مراجع برای استفاده از PJRT در سمت چارچوب، برای ارتباط با دستگاه های PJRT:

• JAX
  • jax-ml/jax با APIهای PJRT از طریق APIهای xla_client تعامل دارد
• GoMLX
  • gomlx/gopjrt
  • gomlx/gomlx > backends/xla
• ZML
  • PJRT API wrapper pjrt.zig
  • بارگذاری افزونه PJRT context.zig
  • تعامل با PJRT Buffers buffer.zig
  • یک ماژول را از طریق PJRT module.zig اجرا کنید

پیاده سازی های سخت افزاری

• پلاگین های ادغام کامل (PJRT+MLIR+XLA):
  • پلاگین XLA CPU
  • پلاگین XLA GPU
  • پلاگین XLA اینتل
• پلاگین های یکپارچه سازی نور (PJRT+MLIR):
  • افزونه StableHLO Reference Interpreter (مبتنی بر MLIR، پلاگین C++، برای پیوند بعد از devlabs)
  • پلاگین Tenstorrent-XLA (بر اساس MLIR، پلاگین C)