Przykłady PJRT

Przykład: wtyczka CUDA JAX

1. Implementacja interfejsu PJRT C API za pomocą obudowy (pjrt_c_api_gpu.h).
2. Skonfiguruj punkt wejścia pakietu (setup.py).
3. Zaimplementuj metodę initialize() (__init__.py).
4. Można go testować za pomocą dowolnych testów jax dla CUDA.

Implementacje platform

Informacje o używaniu PJRT po stronie frameworka w celu nawiązania połączenia z urządzeniami PJRT:

• JAX
  • jax-ml/jax: interfejs API współpracuje z interfejsami PJRT za pomocą interfejsów xla_client
• GoMLX
  • gomlx/gopjrt
  • gomlx/gomlx > backends/xla
• ZML
  • Pakiet API PJRT pjrt.zig
  • Załaduj wtyczkę PJRT context.zig.
  • Interakcje z buforami PJRT buffer.zig
  • Wykonaj moduł za pomocą pakietu module.zig PJRT.

Implementacje sprzętowe

• Wtyczki pełnej integracji (PJRT+MLIR+XLA):
  • Plik wykonywalny XLA dla procesora
  • Wtyczka GPU XLA
  • Wtyczka Intel XLA
• Wtyczki do lekkiej integracji (PJRT+MLIR):
  • StableHLO Reference Interpreter plugin (wtyczka Interpretera referencyjnego StableHLO (oparta na MLIR, wtyczka C++), która zostanie połączona po zakończeniu devlabs).
  • Wtyczka Tenstorrent-XLA (opracowana na podstawie MLIR, wtyczka w języku C)