Örnek: JAX CUDA eklentisi
- Sarmalayıcı aracılığıyla PJRT C API uygulaması (pjrt_c_api_gpu.h).
- Paketin giriş noktasını (setup.py) ayarlayın.
- initialize() yöntemini (__init__.py) uygulayın.
- CUDA için herhangi bir jax testiyle test edilebilir.
Çerçeve Uygulamaları
PJRT cihazlarıyla arayüz oluşturmak için PJRT'yi çerçeve tarafında kullanmayla ilgili bazı referanslar:
- JAX
- jax-ml/jax,
xla_client
API'leri aracılığıyla PJRT API'leriyle etkileşim kurar
- jax-ml/jax,
- GoMLX
- ZML
- PJRT API sarmalayıcısı pjrt.zig
- PJRT eklentisini context.zig yükleme
- PJRT Tamponları buffer.zig ile etkileşim kurma
- PJRT module.zig aracılığıyla modül yürütme
Donanım Uygulamaları
- Tam entegrasyon eklentileri (PJRT+MLIR+XLA):
- Basit entegrasyon eklentileri (PJRT+MLIR):
- StableHLO Referans Yorumlayıcı eklentisi (MLIR tabanlı, C++ eklentisi, devlabs'den sonra bağlanacak)
- Tenstorrent-XLA eklentisi (MLIR tabanlı, C eklentisi)