▼ もうちょっと詳しく

2018年にGoogleが公開。NumPy風の書き味でコードを書きながら、自動微分・GPU/TPU対応・JITコンパイルを全部やってくれる、研究者向けの高速計算ライブラリです。

PyTorchやTensorFlowが「ニューラルネット用」なのに対し、JAXは「もっと低レベルで自由」。研究者には人気で、Google内部の最先端研究では多用されています。手放しでお勧めできるか聞かれると微妙な、玄人寄りの道具。