こんにちは。今回はGoogleの開発するディープラーニング向けライブラリ⭐️であるJAXの紹介をします。 https://github.com/google/jax ディープラーニング向けライブラリとしてはTensorflowやPyTorch、最近開発終了が宣言されたChainerなどが有名かと思います。これらは多次元配列と自動微分をサポートした計算ライブラリをコアとしていて、それにニューラルネットの実装を容易にするラッパーなどが付属しています。 GoogleといえばTensorflowが有名ですが、JAXはTensorflowとは何が違うのでしょうか。 JAXを一言で表現すると、高速なautogradです。 もう少し詳しくいうと、多次元配列の計算ライブラリであるnumpyに自動微分とJITがくっついたものです。さらに、GPUやTPUといったアクセラレーター上でも動作します。Tensorf