Esempi di PJRT

Esempio: plug-in CUDA JAX

1. Implementazione dell'API C PJRT tramite wrapper (pjrt_c_api_gpu.h).
2. Configura il punto di ingresso per il pacchetto (setup.py).
3. Implementa un metodo initialize() (__init__.py).
4. Può essere testato con qualsiasi test jax per CUDA.

Implementazioni dei framework

Alcuni riferimenti per l'utilizzo di PJRT lato framework, per interfacciarsi con i dispositivi PJRT:

• JAX
  • jax-ml/jax interagisce con le API PJRT tramite le API xla_client
• GoMLX
  • gomlx/gopjrt
  • gomlx/gomlx > backends/xla
• ZML
  • Wrapper dell'API PJRT pjrt.zig
  • Carica il plug-in PJRT context.zig
  • Interazione con i buffer PJRT buffer.zig
  • Esegui un modulo tramite module.zig di PJRT

Implementazioni hardware

• Plug-in di integrazione completa (PJRT+MLIR+XLA):
  • Plugin CPU XLA
  • Plugin GPU XLA
  • Plug-in Intel XLA
• Plug-in di integrazione delle luci (PJRT+MLIR):
  • Plug-in di interprete di riferimento StableHLO (plug-in C++ basato su MLIR, da collegare dopo devlabs)
  • Plugin Tenstorrent-XLA (plugin C basato su MLIR)