JAX (ライブラリ)

JAXは、高速な数値計算と大規模な機械学習のために設計されたPythonのオープンソースのライブラリ。 NumPy風の構文で書かれたPythonのソースコードをCPU・GPU・AIアクセラレータへコンパイルする実行時コンパイラや自動微分などを含む。