Rust の std::autodiff モジュールでは、微分可能プログラミングを利用できます。
#![feature(autodiff)]
use std::autodiff::*;
// f(x) = x * x、f'(x) = 2.0 * x
// したがって bar は (x * x, 2.0 * x) を返す
#[autodiff_reverse(bar, Active, Active)]
fn foo(x: f32) -> f32 { x * x }
fn main() {
assert_eq!(bar(3.0, 1.0), (9.0, 6.0));
assert_eq!(bar(4.0, 1.0), (16.0, 8.0));
}std::autodiff モジュールの詳細なドキュメントは std::autodiff で入手できます。
微分可能プログラミングは、数値計算、固体力学、計算化学、流体力学、バックプロパゲーションによるニューラルネットワークのトレーニング、ODE ソルバー、微分可能レンダリング、量子コンピューティング、気候シミュレーションなど、さまざまな分野で使用されています。
std::autodiff は現在、自動微分のための LLVM ベースのツールである Enzyme を基にしています。コンパイラーベースの自動微分に依存する主な理由は 3 つあります。
- 使いやすさ: 現在の自動微分クレートは、通常の Rust プログラムをサポートしていません。カスタム DSL を強制するか、ライブラリが提供する型（たとえばスライスや配列の代わり）を使用する必要があるか、スカラー関数に限定されています。コンパイラーベースの自動微分では、配列、スライス、ユーザー定義の構造体や列挙型、制御フローなどを含む通常の Rust コードをユーザーが記述できます。
- パフォーマンス: 既存の Rust 自動微分アプローチのほとんどには、操作ごとに一定のオーバーヘッドがあります。これは、大きなテンソルに対する高コストな操作が少ない ML アプリケーションでは簡単に償却できます。しかし、HPC や科学技術計算の分野のアプリケーションでは、多くの場合許容できません。（最適化された）LLVM IR に対して動作することで、コンパイラーベースの自動微分は、そのようなケースで大幅に優れたパフォーマンスを実現できます。
- 機能: このような低レベルで動作し、他の LLVM ベース言語と実装を共有することで、Enzyme プロジェクトですでに行われている大量の作業を活用できます。たとえば、MPI ルーチンを呼び出す Rust コードや、CuBLAS のようなライブラリを含む GPU コードをサポートできます。