Exemples de PJRT

Exemple: Plug-in CUDA JAX

1. Implémentation de l'API C PJRT via un wrapper (pjrt_c_api_gpu.h).
2. Configurez le point d'entrée du package (setup.py).
3. Implémentez une méthode initialize() (__init__.py).
4. Peut être testé avec n'importe quel test jax pour CUDA.

Implémentations de frameworks

Quelques références pour utiliser PJRT côté framework, afin d'interagir avec les appareils PJRT:

• JAX
  • jax-ml/jax interagit avec les API PJRT via les API xla_client.
• GoMLX
  • gomlx/gopjrt
  • gomlx/gomlx > backends/xla
• ZML
  • Encapsuleur d'API PJRT pjrt.zig
  • Charger le plug-in PJRT context.zig
  • Interagir avec les tampons PJRT buffer.zig
  • Exécuter un module via PJRT module.zig

Implémentations matérielles

• Plugins d'intégration complète (PJRT+MLIR+XLA) :
  • Plug-in de processeur XLA
  • Plug-in GPU XLA
  • Plug-in Intel XLA
• Plug-ins d'intégration légers (PJRT+MLIR) :
  • Plug-in de l'interprète de référence StableHLO (plug-in C++ basé sur MLIR, à associer après les devlabs)
  • Plug-in Tenstorrent-XLA (plug-in C basé sur le MLIR)