Примеры ПЖРТ

Пример: плагин JAX CUDA.

1. Реализация PJRT C API через оболочку ( pjrt_c_api_gpu.h ).
2. Настройте точку входа для пакета ( setup.py ).
3. Реализуйте метод инициализации() ( __init__.py ).
4. Можно протестировать любыми jax-тестами для CUDA.

Реализации фреймворков

Некоторые ссылки по использованию PJRT на стороне платформы для взаимодействия с устройствами PJRT:

• ДЖАКС
  • jax-ml/jax взаимодействует с API PJRT через API xla_client .
• GoMLX
  • gomlx/gopjrt
  • gomlx/gomlx > бэкенды/xla
• ЗМЛ
  • Оболочка PJRT API pjrt.zig
  • Загрузите плагин PJRT context.zig
  • Взаимодействие с буферами PJRT buffer.zig
  • Запуск модуля через PJRT Module.zig

Аппаратные реализации

• Плагины полной интеграции (PJRT+MLIR+XLA):
  • Плагин XLA для процессора
  • Плагин XLA для графического процессора
  • Плагин Intel XLA
• Плагины легкой интеграции (PJRT+MLIR):
  • Плагин StableHLO Reference Interpreter (плагин C++ на основе MLIR, ссылка будет добавлена ​​после devlabs)
  • Плагин Tenstorrent-XLA (плагин C на основе MLIR)