गूगल जैक्स

JAX

Logo
Developer(s)	Google
Preview release	v0.3.13 / 16 May 2022 (2022-05-16)
Repository	github.com/google/jax
Written in	Python, C++
Operating system	Linux, macOS, Windows
Platform	Python, NumPy
Size	9.0 MB
Type	Machine learning
License	Apache 2.0

Google JAX संख्यात्मक कार्यों को बदलने के लिए एक मशीन लर्निंग फ्रेमवर्क है। इसे autograd (फ़ंक्शन के विभेदन के माध्यम से ग्रेडिएंट फ़ंक्शन की स्वचालित प्राप्ति) और TensorFlow के के एक संशोधित संस्करण को एक साथ लाने के रूप में वर्णित किया गया है। xla XLA (त्वरित रैखिक बीजगणित)। इसे NumPy की संरचना और वर्कफ़्लो का यथासंभव बारीकी से पालन करने के लिए डिज़ाइन किया गया है और यह TensorFlow और PyTorch जैसे विभिन्न मौजूदा फ़्रेमवर्क के साथ काम करता है। JAX के प्राथमिक कार्य हैं:

1. ग्रेड: स्वचालित भेदभाव
2. जिट: संकलन
3. वीमैप: ऑटो-वेक्टराइजेशन
4. pmap: एकल प्रोग्राम, एकाधिक डेटा प्रोग्रामिंग

स्नातक

Main article: Automatic differentiation

नीचे दिया गया कोड दर्शाता है grad फ़ंक्शन का स्वचालित विभेदन।

<सिंटैक्सहाइलाइट लैंग= सुन्न लाइन= 1 >

1. आयात

जैक्स इंपोर्ट ग्रेड से jax.numpy को jnp के रूप में आयात करें

1. लॉजिस्टिक फ़ंक्शन को परिभाषित करें

डीईएफ़ लॉजिस्टिक(x):

   वापसी jnp.exp(x) / (jnp.exp(x) + 1)

1. लॉजिस्टिक फ़ंक्शन का ग्रेडिएंट फ़ंक्शन प्राप्त करें

ग्रैड_लॉजिस्टिक = ग्रैड(लॉजिस्टिक)

1. x = 1 पर लॉजिस्टिक फ़ंक्शन के ग्रेडिएंट का मूल्यांकन करें

grad_log_out = grad_लॉजिस्टिक(1.0) प्रिंट(grad_log_out) </सिंटैक्सहाइलाइट>

अंतिम पंक्ति को आउटपुट करना चाहिएː

सिंटैक्सहाइलाइट लैंग = आउटपुट > 0.19661194 </सिंटैक्सहाइलाइट>

जीत

Main article: Just-in-time compilation

नीचे दिया गया कोड फ़्यूज़न के माध्यम से जिट फ़ंक्शन के अनुकूलन को दर्शाता है।

<सिंटैक्सहाइलाइट लैंग= सुन्न लाइन= 1 >

1. आयात

जैक्स आयात जीआईटी से jax.numpy को jnp के रूप में आयात करें

1. क्यूब फ़ंक्शन को परिभाषित करें

डेफ़ क्यूब(x):

   x * x * x लौटें

1. डेटा जनरेट करें

x = jnp.ones((10000, 10000))

1. क्यूब फ़ंक्शन का जिट संस्करण बनाएं

जिट_क्यूब = जिट(क्यूब)

1. गति तुलना के लिए क्यूब और जिट_क्यूब फ़ंक्शन को समान डेटा पर लागू करें

घन(x) जिट_क्यूब(x) </सिंटैक्सहाइलाइट>

के लिए गणना समय jit_cube (पंक्ति संख्या 17) उससे काफ़ी छोटी होनी चाहिए cube (पंक्ति संख्या 16). लाइन नंबर पर मान बढ़ाना। 7, अंतर बढ़ जाएगा.

वीमैप

Main article: Vectorization (mathematics)

नीचे दिया गया कोड दर्शाता है vmap फ़ंक्शन का वैश्वीकरण।

<सिंटैक्सहाइलाइट लैंग= सुन्न लाइन= 1 >

1. आयात

functools से आंशिक आयात करें जैक्स आयात वीमैप से jax.numpy को jnp के रूप में आयात करें

1. फ़ंक्शन परिभाषित करें

डीईएफ़ ग्रेड (स्वयं, इनपुट):

   in_grad_partial = आंशिक(self._net_grads, self._net_params)
   grad_vmap = vmap(in_grad_partial)
   rich_grads = grad_vmap(इनपुट)
   फ़्लैट_ग्रेड्स = np.asarray(self._flatten_batch(rich_grads))
   फ्लैट_ग्रेड्स.एनडिम == 2 और फ्लैट_ग्रेड्स.शेप[0] == इनपुट्स.शेप[0] पर जोर दें
   फ़्लैट_ग्रैड्स लौटाएँ

</सिंटैक्सहाइलाइट>

इस अनुभाग के दाईं ओर GIF वेक्टरकृत जोड़ की धारणा को दर्शाता है।

सदिशकृत जोड़ का चित्रण वीडियो

pmap

Main article: Automatic parallelization

नीचे दिया गया कोड दर्शाता है pmap मैट्रिक्स गुणन के लिए फ़ंक्शन का समानांतरीकरण।

<सिंटैक्सहाइलाइट लैंग= सुन्न लाइन= 1 >

1. JAX से pmap और रैंडम आयात करें; JAX NumPy आयात करें

जैक्स आयात pmap से, यादृच्छिक jax.numpy को jnp के रूप में आयात करें

1. 5000 x 6000 आयामों के 2 यादृच्छिक मैट्रिक्स उत्पन्न करें, प्रति डिवाइस एक

रैंडम_कीज़ = रैंडम.स्प्लिट(रैंडम.PRNGKey(0), 2) मैट्रिसेस = pmap(लैम्ब्डा कुंजी: रैंडम.नॉर्मल(कुंजी, (5000, 6000)))(random_keys)

1. डेटा ट्रांसफर के बिना, समानांतर में, प्रत्येक सीपीयू/जीपीयू पर एक स्थानीय मैट्रिक्स गुणन करें

आउटपुट = pmap(लैम्ब्डा x: jnp.dot(x, x.T))(मैट्रिसेस)

1. डेटा ट्रांसफर के बिना, समानांतर में, प्रत्येक सीपीयू/जीपीयू पर दोनों मैट्रिक्स के लिए अलग-अलग माध्य प्राप्त करें

मतलब = pmap(jnp.mean)(आउटपुट) प्रिंट(मतलब) </सिंटैक्सहाइलाइट>

अंतिम पंक्ति को मान मुद्रित करना चाहिएː

सिंटैक्सहाइलाइट लैंग = आउटपुट > [1.1566595 1.1805978] </सिंटैक्सहाइलाइट>

JAX का उपयोग करने वाली लाइब्रेरी

कई पायथन लाइब्रेरीज़ JAX को बैकएंड के रूप में उपयोग करती हैं, जिनमें शामिल हैं:

• फ़्लैक्स, एक उच्च स्तरीय न्यूरल नेटवर्क लाइब्रेरी जिसे प्रारंभ में Google Brain द्वारा विकसित किया गया था।
• इक्विनॉक्स, एक लाइब्रेरी जो पैरामीटरयुक्त कार्यों (न्यूरल नेटवर्क सहित) को PyTrees के रूप में प्रस्तुत करने के विचार के इर्द-गिर्द घूमती है। इसे पैट्रिक किडगर ने बनाया था।
• डिफ्रैक्स, साधारण अंतर समीकरणों के लिए संख्यात्मक तरीकों के लिए एक पुस्तकालय, जैसे साधारण अंतर समीकरण और स्टोचैस्टिक अंतर समीकरण।
• ऑप्टैक्स, डीपमाइंड द्वारा विकसित ग्रेडिएंट प्रोसेसिंग और गणितीय अनुकूलन के लिए एक पुस्तकालय।
• Lineax, रैखिक समीकरणों की प्रणाली और रैखिक न्यूनतम वर्गों के लिए संख्यात्मक तरीकों के लिए एक पुस्तकालय।
• RLax, डीपमाइंड द्वारा विकसित सुदृढीकरण शिक्षण एजेंटों को विकसित करने के लिए एक पुस्तकालय।
• जेराफ, ग्राफ न्यूरल नेटवर्क के लिए एक लाइब्रेरी, जिसे डीपमाइंड द्वारा विकसित किया गया है।
• जैक्सटाइपिंग, टाइप सिग्नेचर जोड़ने के लिए एक लाइब्रेरी सरणियों या टेंसरों के आकार और डेटा प्रकार (dtype) के लिए।

कुछ R (प्रोग्रामिंग भाषा) लाइब्रेरी JAX को बैकएंड के रूप में भी उपयोग करती हैं, जिनमें शामिल हैं:

• fastrerandomize, एक लाइब्रेरी जो Rerandomization के रूप में ज्ञात प्रयोग प्रक्रिया के डिजाइन में संतुलित रैंडमाइजेशन के चयन में तेजी लाने के लिए JAX में रैखिक-बीजगणित अनुकूलित कंपाइलर का उपयोग करती है।

यह भी देखें

• नम्पी
• टेंसरफ्लो
• पाइटोरच
• क्यूडा
• स्वचालित भेदभाव
• सही समय पर संकलन
• वैश्वीकरण (बहुविकल्पी)
• स्वचालित समानांतरीकरण

बाहरी संबंध

• Documentationː jax.readthedocs.io
• Colab (Jupyter/IPython) Quickstart Guideː colab.research.google.com/github/google/jax/blob/main/docs/notebooks/quickstart.ipynb
• TensorFlow's XLAː www.tensorflow.org/xla (Accelerated Linear Algebra)
• Intro to JAX: Accelerating Machine Learning research on YouTube
• Original paperː mlsys.org/Conferences/doc/2018/146.pdf