JuliaからRustへ：科学計算のための微分可能なテンソルスタック

JuliaからRustへ：エージェントAI時代の科学計算のための微分可能なテンソルスタック。 tenferro-rsはRustネイティブの密テンソルスタックです：線形代数、PyTorchスタイルの即時自動微分、JAXスタイルのトレース変換、NumPyスタイルのアインシュタイン summation、FFT、拡張可能な操作クレート、そして明示的なCPU/CUDAバックエンドを備えています。最初のクレートは2026年6月23日（JST）にcrates.ioで公開されました。 著：廣志品岡（埼玉大学）、tensor4allチーム 🌐英語・日本語・簡体中文 ほとんどのテンソルネットワークコードはJuliaで書かれており、私たちも例外ではありませんでした。ITensorsとその周辺エコシステムはプロトタイピングに適しています。コードは数学に近く、反復が容易です。IR基底、スパースモデリング、そしてtensor4allのテンソル交差補間および量子スタックに関する私たち自身の研究もそこから始まりました。しかし、コードベースが大きくなると、Juliaでの開発は遅くなり始めます。実行時にのみ現れる型不安定性、編集/テストループを長くするコンパイル時間とプリコンパイル時間、そしてコードが大きくなるにつれて正確性の確認が難しくなる感覚です。テンソルネットワークスタックをより大きなシステムに組み込み始めたとき、これは無視できない問題となりました。そこで私たちは計算エンジンをRustに移行し始めました。これがすぐに第二の問題を露呈しました。私たちが構築しようとしていたテンソルライブラリはまだ存在しなかったのです。Rustには個々のタスクのためのライブラリはあります。配列のためのndarray、ディープラーニングのためのBurn、線形代数のためのfaerなどです。しかし、アインシュタイン summationを介した自動微分をカバーし、科学計算にも利用できるテンソル層が欠けていました。目標はこれらのライブラリを置き換えることではなく、既存のピースを接続することでした。 Rustエコシステムはここ数年で大きく変化しました。crates.ioは2015年の602クレートから2026年には約21万クレートに増加しました（データ）。密な線形代数にはfaer、GPUカーネルにはCubeCL、一般的な数値計算にはnum-traitsとnum-complexがあります。また、周辺の層にもライブラリがあります。配列のためのndarray、線形代数のためのnalgebraとfaer、ディープラーニングのためのBurnとcandle、NumPyスタイルの配列APIのためのnumrなどです。私たちが必要としていたのは、それらの間の層でした。つまり、列優先ストレージ、動的シェイプ、即時およびトレース自動微分、アインシュタイン summation、FFT、CPU/CUDAバックエンド、そして拡張可能な操作を備えた科学計算テンソルスタックです。それがtenferro-rsの目的です。私たちはfaerとCubeCLを基盤とし、不足している部分を再発明するのではなく追加しています。SparseIR.jlとJuliaテンソルネットワークコードの移植によって、その不足している層がどこにあるのかがより明確になりました。それがtenferro-rsの背景です。この投稿では、なぜ私たちがそれを構築しているのか、そしてコードが人間だけによって書かれる時代ではない今、なぜRustを選んだのかを説明します。 以前はJuliaで問題なかったのに、なぜ今Rustなのか？ 数年前なら、私は学生にJuliaから始めるように言ったでしょう。Juliaコードは数学に近く、メモリ管理は簡単で、数値ライブラリもすでに揃っていました。Rustは学ぶことが多く、エコシステムにはまだ不足している部分がありました。しかし、今では同じアドバイスはしません。Rustが変わったからではなく、コードのほとんどを書くのが私ではなくなったからです。Fortran、Python、Juliaはすべて、人間が手作業でコードを書き、読み、保守するコストを下げるために開発されました。読みやすさ、REPL、数学に近い記法、低い参入障壁がすべて重要でした。AIがコードの大部分を書くようになると、トレードオフが変わります。記述速度はそれほど重要ではなくなります。学習コストの多くはエージェントが処理できます。しかし、「数学のように読める」からといって正確性が保証されるわけではありません。エイリアシング、ミューテーション、アロケーションはコードの表面からは見えません。私たちにとっての問いは、「人間がこれをどれだけ速く書けるか？」ではなく、「これが正しいとどれだけ確信できるか？」に変わりました。この再構築こそが、Rustがより実用的な選択肢となった理由です。具体的には： 所有権と型はコンパイル時に広範なエラーを排除します。cargo checkは数秒で回答を返すので、エージェントが何かを間違えても、プログラムを実行する前に発見できます。 Cargoはビルド、依存関係解決、テスト、ベンチマークを1か所で処理します。CMakeは不要で、リンク時のバージョン競合もありません。完全なスタックと依存関係をラップトップでゼロからビルドしても数分で完了し、編集/テストループは数十秒です。 Rustはモジュールとクレートの境界に沿ってシンボル可視性を制御します。エージェントは層内でのみ機能し、別のクレートの内部に踏み込んで抽象化を密かに破壊することはありません。約13万行のAI生成コードベースでは、この境界が重要です。 ライフタイムと所有権のメカニズムは大部分をエージェントに任せることができるため、人間の注意はアルゴリズム、設計、正確性に向けられます。かつてRustの欠点とされていた初期学習コストは、今ではそれほど問題ではありません。C++、Python、Juliaでは、大規模なコードベースは検証が難しくなりすぎると懸念される傾向がありますが、Rustではその懸念は著しく小さくなります。 ポートからスタックへ 私たちは一般的なテンソルライブラリを構築するつもりはありませんでした。必要なピースを移植し、ツールとの格闘に費やす時間を減らしたかったのです。しかし、実装が進むにつれて、自動微分、バックエンド、そして新しい操作が追加される方法は、テンソルネットワーク固有の層の中に閉じ込められるべきではないことが明らかになりました。私たちは共有部分を独立したテンソルスタックとして設計しました。操作ファミリーは、オールインワンのテンソル型の中にではなく、独自のクレートに存在します。自動微分ルールはテンソル型の外に存在します。Julia/ChainRulesの教訓に従い、導関数ルールは特定の具象テンソルクラスではなく、操作自体に属します。AD基盤であるtidu-rsはジェネリックであり、テンソル型はそのコンシューマの1つに過ぎません。バックエンドとデバイスは明示的です。CPUとGPU間でデータが密かに移動することはありません。また、どのバックエンドが操作を実行できるかという問題と、実行時にどのデバイスが利用可能かという問題を分離しています。ストレージは列優先であり、Fortran、Julia、MATLAB、LAPACK/BLASに一致します。行優先データも、不必要な即時コピーなしにストライドビューを介して処理できます。この設計により、スタックはテンソルネットワーク以外でも役立ちます。 tenferro-rsを2分で紹介 このスタックは、型付きテンソル、backward()による即時（イーガー）実行、grad/vjp/jvp/HVPによるトレースグラフ、線形代数、アインシュタイン summation、FFT、および明示的なCPUとCUDAバックエンド（および実験的なWebGPU）を提供します。以下は、sum(x²)の勾配が2xであるPyTorchスタイルの即時自動微分を、リポジトリのeager_autodiff_pytorch_style.rsからそのままコピーしたものです。 ```rust use tenferro_ad::{EagerRuntime, Tensor}; fn assert_close(actual: &[f64], expected: &[f64]) { assert_eq!(actual.len(), expected.len()); for (index, (actual, expected)) in actual.iter().zip(expected).enumerate() { let error = (actual - expected).abs(); assert!( error < 1.0e-12, "value {index}: actual={actual}, expected={expected}, error={error}" ); } } fn main() -> Result<(), Box<dyn std::error::Error>> { let runtime = EagerRuntime::new(); let x = runtime.variable_from(Tensor::from_vec_col_major( vec![3], vec![1.0_f64, 2.0, 3.0], )?)?; let prediction = x.mul(&x).unwrap(); let loss = prediction.reduce_sum(&[0])?; assert_eq!(loss.shape(), &[]); assert_close(loss.materialized()?.as_slice::<f64>().unwrap(), &[14.0]); loss.backward()?; let grad = x .grad()?; .expect("tracked variable should receive a gradient"); assert_eq!(grad.shape(), &[3]); assert_close(grad.as_slice::<f64>().unwrap(), &[2.0, 4.0, 6.0]); x.clear_grad()?; assert!(x.grad()?.is_none()); Ok(()) } ``` この例では、プログラムが自身の結果をチェックしています。この小さな例は、私たちがライブラリを構築する方法を反映しています。同じ計算は、JAXスタイルのトレースグラフとしても実行でき、一度コンパイルされて再利用され、grad/vjp/jvpを伴います（traced_autodiff_jax_style.rsを参照）。必要な層を選択できます。型付きテンソル、自動微分による即時実行、またはトレースグラフです。自動微分、CUDA、アインシュタイン summation、FFT、および線形代数は、必要な場合にのみ有効にできます。サイズがデータに依存する計算。JAXとXLAは、形状が固定された計算を最適化し、同じ形状の入力に対して迅速に再実行することに優れています。形状がデータに依存するようになると、事態は難しくなります。 ---