PJRT-Beispiele

Beispiel: JAX-CUDA-Plug-in

1. PJRT C API-Implementierung über Wrapper (pjrt_c_api_gpu.h)
2. Richten Sie den Einstiegspunkt für das Paket (setup.py) ein.
3. Implementieren Sie die Methode „initialize()“ (__init__.py).
4. Kann mit allen jax-Tests für CUDA getestet werden.

Framework-Implementierungen

Einige Referenzen zur Verwendung von PJRT auf Framework-Seite, um mit PJRT-Geräten zu kommunizieren:

• JAX
  • jax-ml/jax interagiert über die xla_client APIs mit PJRT APIs
• GoMLX
  • gomlx/gopjrt
  • gomlx/gomlx > backends/xla
• ZML
  • PJRT API-Wrapper pjrt.zig
  • PJRT-Plug-in context.zig laden
  • Interaktion mit PJRT-Bussen buffer.zig
  • Modul über PJRT module.zig ausführen

Hardwareimplementierungen

• Plug-ins für die vollständige Integration (PJRT+MLIR+XLA):
  • XLA-CPU-Plug-in
  • XLA-GPU-Plug-in
  • Intel XLA-Plug-in
• Plug-ins für die einfache Integration (PJRT+MLIR):
  • StableHLO Reference Interpreter-Plug-in (MLIR-basiertes C++-Plug-in, das nach dem DevLab verknüpft werden muss)
  • Tenstorrent-XLA-Plug-in (MLIR-basiert, C-Plug-in)