Google Jax - उच्च-प्रदर्शन मशीन लर्निंग लाइब्रेरी

"आसानी से इस्तेमाल होने वाली चीज़ों को तेज़ बनाकर मशीन लर्निंग में लागू करना"

• केवल Python और Numpy का संयोजन

→ XLA का उपयोग करके Numpy को GPU/TPU पर compile और run करना

→ Python functions को एक ही API से JIT compile करके XLA-optimized kernels में आसानी से डाला जा सकता है

→ कई GPU/TPU पर execution भी आसान (vmap, pmap)

• मौजूदा Python+Numpy प्रदर्शन से कहीं आगे

[xguru]

DeepMind ने पूरी चीज़ को Jax-आधारित तरीके से refactor किया है

https://deepmind.com/blog/article/using-jax-to-accelerate-our-research