इस प्रोजेक्ट के बारे में
JAX, Python और NumPy प्रोग्राम्स के कंपोजेबल फंक्शन ट्रांसफॉर्मेशन के लिए एक सिस्टम प्रदान करता है। इसे संकलन (compilation) के लिए XLA (Accelerated Linear Algebra) का उपयोग करके GPUs और TPUs सहित हार्डवेयर एक्सेलेरेटर्स पर न्यूमेरिकल कंप्यूटिंग को स्केल करने के लिए डिज़ाइन किया गया है।
मुख्य क्षमताओं में शामिल हैं:
- ऑटोमैटिक डिफरेंशिएशन: `jax.grad` का उपयोग करते हुए, यह Python लूप्स, ब्रांचेस, रिकर्शन और क्लोजर्स के माध्यम से रिवर्स-मोड और फॉरवर्ड-मोड डिफरेंशिएशन का समर्थन करता है, जिससे हायर-ऑर्डर डेरिवेटिव्स संभव होते हैं।
- जस्ट-इन-टाइम (JIT) कंपाइलेशन: `jax.jit` का उपयोग करते हुए, यह एक्सेलेरेटर्स पर निष्पादन प्रदर्शन (execution performance) में सुधार के लिए XLA के माध्यम से प्योर फंक्शन्स को कंपाइल करता है।
- ऑटो-वेक्टराइजेशन: `jax.vmap` का उपयोग करते हुए, यह लूप्स को प्रिमिटिव ऑपरेशन्स तक धकेल कर एरे एक्सिस के साथ फंक्शन्स को मैप करता है, जिससे मैनुअल बैच डायमेंशन मैनेजमेंट की आवश्यकता समाप्त हो जाती है।
- स्केलिंग और पैरेललाइजेशन: यह विभिन्न स्केलिंग मोड का समर्थन करता है जिसमें कंपाइलर-आधारित ऑटोमैटिक पैरेललाइजेशन, एक्सप्लिसिट शार्डिंग और एक्सप्लिसिट कलेक्टिव्स के साथ मैनुअल प्रति-डिवाइस प्रोग्रामिंग शामिल है।
JAX, Linux, macOS और Windows सहित कई प्लेटफॉर्म्स के साथ संगत है, जिसमें CPU, NVIDIA GPU, Google TPU, AMD GPU और Apple एवं Intel GPUs के लिए प्रयोगात्मक (experimental) समर्थन शामिल है।