jax

JAXは、PythonおよびNumPyプログラムの合成可能な関数変換システムを提供します。XLA (Accelerated Linear Algebra) によるコンパイルを用いて、GPUやTPUを含むハードウェアアクセラレータ全体で数値計算をスケールさせるように設計されています。 主な機能は以下の通りです： - 自動微分：`jax.grad` を使用し、Pythonのループ、分岐、再帰、クロージャを通じたリバースモードおよびフォワードモードの微分をサポートし、高階導関数を可能にします。 - Just-In-Time (JIT) コンパイル：`jax.jit` を使用し、純粋関数をXLA経由でコンパイルすることで、アクセラレータ上の実行パフォーマンスを向上させます。 - 自動ベクトル化：`jax.vmap` を使用し、ループをプリミティブ操作に押し下げることで関数を配列軸に沿ってマッピングし、手動でのバッチ次元管理を不要にします。 - スケーリングと並列化：コンパイラベースの自動並列化、明示的なシャーディング、および明示的なコレクティブを用いた手動のデバイス別プログラミングを含む、さまざまなスケーリングモードをサポートしています。 JAXはLinux、macOS、Windowsを含む複数のプラットフォームと互換性があり、CPU、NVIDIA GPU、Google TPU、AMD GPUをサポートしているほか、AppleおよびIntel GPUへの実験的なサポートを提供しています。