गूगल जैक्स

जैक

Logo
Developer(s)	गूगल
Preview release	v0.3.13 / 16 May 2022 (2022-05-16)
Repository	github.com/google/jax
Written in	पायथन, सी ++
Operating system	लिनक्स, मैकओएस, विंडोज
Platform	पायथन, नमपाई
Size	9.0 MB
Type	मसीन लर्निंग
License	अपाचे 2.0

गूगल जैक्स न्यूमेरिकल फ़ंक्शन को परिवर्तित करने के लिए एक प्रकार का मशीन लर्निंग फ्रेमवर्क है। इसे ऑटोग्रैड के माध्यम से ग्रेडिएंट फ़ंक्शन और टेंसरफ्लो फ्रेमवर्क के एक्सएलए फ़ंक्शन के एक संशोधित संस्करण को एक साथ प्रस्तुत करने के लिए वर्णित किया गया है। सामान्यतः इसको NumPy के स्ट्रक्चर और वर्कफ़्लो के ग्रेडिएंट फ़ंक्शन का उपयोग करने के लिए डिज़ाइन किया गया है। इसमे सम्मिलित विभिन्न फ्रेमवर्क टेंसरफ्लो और PyTorch एक साथ कार्य करते है। जेएएक्स फ़ंक्शन के प्राथमिक संस्करण हैं:

1. ग्रैड: ऑटोडिफरेंसीएसन
2. जिट: कंपाइलेशन
3. वीमैप: ऑटोवेक्टराइजेशन
4. पीमैप: एसपीएमडी डेटा प्रोग्रामिंग

ग्रैड

Main article: ऑटोडिफरेंसीएसन

नीचे दिया गया कोड grad फ़ंक्शन के ऑटोडिफरेंसीएसन को प्रदर्शित करता है:

# imports
from jax import grad
import jax.numpy as jnp

# define the logistic function
def logistic(x):  
    return jnp.exp(x) / (jnp.exp(x) + 1)

# obtain the gradient function of the logistic function
grad_logistic = grad(logistic)

# evaluate the gradient of the logistic function at x = 1 
grad_log_out = grad_logistic(1.0)   
print(grad_log_out)

0.19661194

जिट

Main article: कंपाइलेशन

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

# imports
from jax import jit
import jax.numpy as jnp

# define the cube function
def cube(x):
    return x * x * x

# generate data
x = jnp.ones((10000, 10000))

# create the jit version of the cube function
jit_cube = jit(cube)

# apply the cube and jit_cube functions to the same data for speed comparison
cube(x)
jit_cube(x)

jit_cube फ़ंक्शन (पंक्ति संख्या 17) के लिए कंपाइलर टाइम cube फ़ंक्शन (पंक्ति संख्या 16) की तुलना में अपेक्षाकृत कम होना चाहिए, क्योंकि पंक्ति संख्या 7 पर मान बढ़ने से अंतर बढ़ सकता है।

वीमैप

Main article: वेक्टराइजेशन

नीचे दिया गया कोड vmap फ़ंक्शन के वेक्टराइजेशन को प्रदर्शित करता है:

# imports
from functools import partial
from jax import vmap
import jax.numpy as jnp

# define function
def grads(self, inputs):
    in_grad_partial = partial(self._net_grads, self._net_params)
    grad_vmap = vmap(in_grad_partial)
    rich_grads = grad_vmap(inputs)
    flat_grads = np.asarray(self._flatten_batch(rich_grads))
    assert flat_grads.ndim == 2 and flat_grads.shape[0] == inputs.shape[0]
    return flat_grads

प्रायः यह फ़ंक्शन जीआईएफ वेक्टराइज्ड संबंध की धारणा को दर्शाता है।

जीआईएफ वेक्टराइज्ड संबंध

पीमैप

Main article: ऑटोमैटिक पैरलाइजेसन

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

# import pmap and random from JAX; import JAX NumPy
from jax import pmap, random
import jax.numpy as jnp

# generate 2 random matrices of dimensions 5000 x 6000, one per device
random_keys = random.split(random.PRNGKey(0), 2)
matrices = pmap(lambda key: random.normal(key, (5000, 6000)))(random_keys)

# without data transfer, in parallel, perform a local matrix multiplication on each CPU/GPU
outputs = pmap(lambda x: jnp.dot(x, x.T))(matrices)

# without data transfer, in parallel, obtain the mean for both matrices on each CPU/GPU separately
means = pmap(jnp.mean)(outputs)
print(means)

[1.1566595 1.1805978]

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

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

• फ्लेक्स, यह एक प्रकार की न्यूरल नेटवर्क लाइब्रेरी है जिसे गूगल ब्रेन द्वारा विकसित किया गया था।
• इक्विनॉक्स, एक लाइब्रेरी जो पैरामीटरयुक्त फ़ंक्शन (न्यूरल नेटवर्क सहित) को PyTrees के रूप में प्रस्तुत करती है।
• डिडिफ्रेक्स, साधारण डिफरेंशियल फ़ंक्शन और स्टोकेस्टिक डिफरेंशियल फ़ंक्शन जैसे फ़ंक्शन के संख्यात्मक समाधान के लिए लाइब्रेरी।
• ऑप्टैक्स, डीपमाइंड द्वारा विकसित ग्रेडिएंट प्रोसेसिंग और ऑप्टिमाइज़ेशन के लिए लाइब्रेरी।
• लाइनेक्स, संख्यात्मक रूप से लीनियर सिस्टम को हल करने के लिए लाइब्रेरी।
• लक्स, डीपमाइंड द्वारा विकसित रिइंफोर्समेंट लर्निंग के एजेंटों को विकसित करने के लिए लाइब्रेरी।
• जेराफ, ग्राफ न्यूरल नेटवर्क के लिए एक लाइब्रेरी, जिसे डीपमाइंड द्वारा विकसित किया गया है।
• जैक्सटाइपिंग, टेबल या टेन्सर्स के आकार और डेटा टाइप ("डीटाइप") के लिए टाइप एनोटेशन इन्सर्ट करने के लिए लाइब्रेरी।

कुछ आर (प्रोग्रामिंग लैंग्वेज) लाइब्रेरी जेएएक्स को बैकएंड के रूप में भी उपयोग करती हैं, जिनमें प्रायः फास्टरेरैंडमाइज़ सम्मिलित हैं:

• फास्टरेरैंडमाइज़, सामान्यतः यह लाइब्रेरी जेएएक्स में लीनियर ऑप्टिमाइजेशन कंपाइलर का उपयोग करती है जिससे परीक्षण प्रक्रिया के डिजाइन में रैंडमाइजेशन के चयन को अधिक तीव्रता से किया जा सकता है, जिसके कारण इसे रीरैंडमाइजेशन के रूप में भी जाना जाता है।

यह भी देखें

• टेंसरफ्लो
• पाइटॉर्च
• क्यूडा
• ऑटोमैटिक डिफरेंसीएसन
• कंपाइलेशन
• वेक्टराइजेशन
• एक्सेलरेटेड लीनियर एल्जेब्रा (रैखिक बीजगणित)

बाहरी संबंध

• 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