Exemplos de PJRT

Exemplo: plug-in CUDA do JAX

1. Implementação da API PJRT C por wrapper (pjrt_c_api_gpu.h).
2. Configure o ponto de entrada do pacote (setup.py).
3. Implemente um método initialize() (__init__.py).
4. Pode ser testado com qualquer teste jax para CUDA.

Implementações de frameworks

Algumas referências para usar o PJRT no framework para interagir com dispositivos PJRT:

• JAX
  • jax-ml/jax interage com as APIs PJRT pelas APIs xla_client
• GoMLX
  • gomlx/gopjrt
  • gomlx/gomlx > backends/xla
• ZML
  • Wrapper da API PJRT pjrt.zig
  • Carregar o plug-in PJRT context.zig
  • Como interagir com buffers PJRT buffer.zig
  • Executar um módulo pelo module.zig do PJRT

Implementações de hardware

• Plug-ins de integração total (PJRT+MLIR+XLA):
  • Plug-in de CPU do XLA
  • Plug-in de GPU XLA
  • Plug-in XLA da Intel
• Plug-ins de integração leve (PJRT+MLIR):
  • Plug-in de intérprete de referência StableHLO (plug-in C++ baseado em MLIR, a ser vinculado após os devlabs)
  • Plug-in tenstorrent-XLA (baseado em MLIR, plug-in C)