PJRT के उदाहरण

उदाहरण: JAX CUDA प्लग इन

1. रैपर (pjrt_c_api_gpu.h) की मदद से PJRT C API लागू करना.
2. पैकेज के लिए एंट्री पॉइंट (setup.py) सेट अप करें.
3. initialize() तरीका (__init__.py) लागू करें.
4. CUDA के लिए, किसी भी jax टेस्ट के साथ जांच की जा सकती है.

फ़्रेमवर्क लागू करना

PJRT डिवाइसों के साथ इंटरफ़ेस करने के लिए, फ़्रेमवर्क साइड पर PJRT का इस्तेमाल करने के कुछ रेफ़रंस:

• JAX
  • jax-ml/jax, xla_client एपीआई के ज़रिए PJRT एपीआई के साथ इंटरैक्ट करता है
• GoMLX
  • gomlx/gopjrt
  • gomlx/gomlx > backends/xla
• ZML
  • PJRT API रैपर pjrt.zig
  • PJRT प्लगिन context.zig लोड करें
  • PJRT बफ़र buffer.zig के साथ इंटरैक्ट करना
  • PJRT module.zig की मदद से मॉड्यूल को चलाना

हार्डवेयर लागू करना

• पूरे इंटिग्रेशन वाले प्लग इन (PJRT+MLIR+XLA):
  • XLA सीपीयू प्लग इन
  • XLA GPU प्लग इन
  • Intel XLA प्लग इन
• लाइट इंटिग्रेशन प्लग इन (PJRT+MLIR):
  • StableHLO रेफ़रंस इंटरप्रेटर प्लग इन (MLIR पर आधारित, C++ प्लग इन, जिसे devlabs के बाद लिंक किया जाएगा)
  • Tenstorrent-XLA प्लग इन (MLIR पर आधारित, C प्लग इन)