patrick-kidger/diffrax

What it solves

Diffrax は、JAX エコシステム内で微分方程式（特に常微分方程式 (ODE)、確率微分方程式 (SDE)、制御微分方程式 (CDE)）を数値的に解く手段を提供します。これらのソルバーは自動微分可能で GPU 上で実行でき、物理ベースのモデルと機械学習の統合に不可欠です。

How it works

JAX 上に構築された本ライブラリは、統一された内部構造を用いてさまざまなタイプの方程式を扱います。Tsit5 や Dopri8 といった多数のソルバーをサポートし、状態を PyTree として表現でき、バックプロパゲーション用の複数の adjoint 手法を提供します。JAX ベースであるため、積分領域を含むプロセスのほぼすべてが vmappable です。

Who it’s for

科学計算や AI に従事する研究者・開発者、特にニューラル微分方程式を実装する方や高性能で微分可能な数値ソルバーが必要な方を対象としています。

Highlights

• ODE、SDE、CDE ソルバーをサポート。
• 完全に自動微分可能で GPU 互換。
• ライブラリ全体で vmappable な操作。
• ニューラル微分方程式をサポート。
• シンプレクティックや暗黙的オプションを含む多様なソルバーを搭載。