グーグルのクラウドを支えるテクノロジー　＞　第140回　オリジナル論文から学ぶ「JAX」の特徴とその役割

＃＃＃

はじめに
　今回は、2018年に公開された論文「Compiling machine learning programs via high-level tracing」を元にして、機械学習ライブラリーJAXの役割と、その基礎となる考え方を紹介します。

JAXのオリジナル論文
　最近、機械学習に関連するオープソースとして、JAXの名前を耳にすることが増えてきました。たとえば、2022年6月に公開されたGoogle Cloudの公式ブログ「EvoJAX: あなたの課題をNeuroevolutionの力で解く」の冒頭には、次のような一節があります。

『JAXはユーザーコードの簡略化や大規模な並列化・何桁もの高速化を可能にする、最近のGoogleで最も重要な機械学習（ML）フレームワークの一つです。このフレームワークは、言語理解におけるPathways Language Model（PaLM） 、物理学・分子動力学シミュレーションにおけるBraxやJAX MDなどをはじめとして、近年に最先端（State-of-the-Art）の成果を示した研究でも利用されています。』

この続きは以下をご覧ください
リンク