Zig で自動微分ライブラリ作ってみた

最近 Bun で注目を浴びている Zig ですが、低レベルシステムプログラミング言語であり、特にメモリ管理周りのカスタマイズが細かいところに手が届く代物のようです。

https://ziglang.org/

しかし、実際に役に立つかどうかを確かめるには、何かに役立つソフトウェアを実装してみるのが一番確実な方法です。そんなわけで、最近はまっている自動微分を実装してみました。

https://github.com/msakuta/zigrad

基本的には rustograd で実装した Tape の再現で、フォワードモード、リバースモード、サブグラフによる高階微分、 graphviz による計算グラフの可視化などの機能を備えますが、スカラーの f64 のみを変数型としてサポートしています。

ただし、 Zig にはまだ公式のパッケージマネージャが存在せず、ライブラリとして利用できる形になっているかはよくわかりません。

出力例

• $sin(x^2)$ の微分

• 高階微分

• グラフの可視化

感想

実際に使ってみた感想です。筆者は Rust のヘビーユーザーなので相応のバイアスがかかっている点についてはご容赦ください。

アロケータ

まず、メモリアロケータに関する自由度が高いのは確かです。 zigrad のデモではアリーナアロケータを使って、 Tape が使うメモリをまとめて管理し、一度に解放しています。個々のメモリを個別に解放するよりも高速であることが期待できます。また、ライブラリ自身はアロケータを引数に取ることによって指定できるので、どのようなアロケータとも組み合わせられます。また、 Go 言語のような defer 文を備えており、スコープから抜けるときに確実にリソースを開放することができます。

var aa = std.heap.ArenaAllocator.init(std.heap.page_allocator);
defer aa.deinit();
var allocator = aa.allocator();

var tape = try Tape.new(&allocator, 32);

しかし、実運用上は常に快適とは言い難いのも否定できません。例えば、次に示すのは掛け算の微分からサブグラフを生成する部分の例ですが、それぞれの演算全てにアロケータの引数を渡さなくてはならず、さらにメモリ確保に失敗する可能性もあるので try を前置する必要があります。これはなかなかの視覚的ノイズです。少なくとも Rust のような ? 後置演算子であれば括弧の数は減らせるのですが。

try (try lhs.mul(drhs, allocator)).add(try dlhs.mul(rhs, allocator), allocator)

また、ライブラリがアロケータと完全に独立に実装できるわけでもなく、 Tape は内部で文字列の確保にアロケータを使っていますが、それらを明示的に解放しないので、呼び出し側のアロケータのスタックのどこかにアリーナが存在していることを想定しています。そうでなければメモリーリークになります。これは型システムでは表現できないのでライブラリの利用者が理解して気を付ける必要があります。

それでも、 Rust のアロケータプロトコルは安定化されていないので、 Zig の利点と言えるかもしれません(Zig はそもそも 1.0 に達しておらず不安定ですが)。

comptime

Zig ではジェネリック型が無い代わりに、型をコンパイル時変数として扱うことができます。例えば、複数のスライスを連結して一続きのメモリ領域を確保する std.mem.concatWithSentinel という標準ライブラリ関数がありますが、その第２引数はスライスの要素型で、ジェネリック関数と同じように使えます。 Sentinel というのはスライスの終端を特殊な値で表すもので、 C の文字列のヌル文字 '\0' に当たります。

ただし、型推論の対象にはならないので、毎回型の名前を書かなくてはならず、高度なジェネリック型や関数の実装には限界があると思われます。

また、これは個人的な感覚ですが、型引数と通常の引数(Zig の用語ではコンパイル時変数と実行時変数)は扱いが分かれていた方が理解しやすいような気がします。 C++ や Rust や最近の Java では、山括弧 <> が型引数の囲み記号として使われており、ジェネリック関数やジェネリック型の型パラメータがどれだか一目でわかります。これが普通の引数に混ざってしまうと認知負荷が増えるように思います。

実際の呼び出し例は次のような感じです。第2引数が u8 という型引数になっているのですが、コードに埋もれている感じがします。

try std.mem.concatWithSentinel(arena.*, u8, &[_][:0]const u8{ "-", tape.*.nodes[self.*.idx].name }, 0)

エラー処理

Zig は Rust と似たような直和型(Sum type)を使ってエラーや Optional を表現しています。ただし、 Rust よりも言語仕様に密接に結びついており、エラー型専用の構文 ! や Optional 専用の ? があります。一般的には専用の構文があったほうが簡潔にコードが書けるといえますが、若干一貫性に欠ける印象があります。例えばエラー型の中に Optional 型を含む型を返す関数のシグネチャは次のようになります。

pub fn gen_graph(...) !?TapeTerm

これを使う側は try と orelse を組み合わせてエラー処理することができますが、どちらが先に適用されるのかコード上だけではわかりません。

const derived = try gaussian.gen_graph(x, &allocator) orelse @panic("derived var doesn't exist");

これに対し、 Rust のメソッドチェイニングであればネストされた型の順番は自明です。

let derived = gaussian.gen_graph(x)?.expect("derived var doesn't exist");

また、とりえあずエラーや null の可能性は無視してコンパイルを通したい場合、エラー型と Optional 型で方法が異なり、 Rust の unwrap() のような覚えやすい方法がないのも気になりました。とは言え、ほとんどが慣れの問題で、しばらく使っていれば気にならなくなるでしょう。

さらに、エラーの種類は一種の集合型になっており、型宣言の ! の前に明示することもできますが、省略すれば関数内で生じる全てのエラーの集合を自動的に返り値の型にしてくれるという機能もあります。しかし、実運用上はエラー型の指定が面倒で省略してしまう場合がほとんどではないかと思います。この点、 Rust ではエラーの集合を明示的に指定するか、 Result<_, Box<dyn Error>> を使うことになり、型消去されていることが明確になり、エラー処理をまじめに考えるようにプログラマを奨励する効果があるかと思います。

Duck typing

個人的には一番気になったのがこれです。 Zig では Duck typing が使われており、コンパイル時型引数を使ったジェネリックなロジックはコンパイルしてみるまでエラーがあっても報告されません。インターフェースやトレイトで型が満たすべき制約を指定することもできません。これは C++ のテンプレートと同じ問題で、 C++ がコンセプトを後から実装したことから推測できるように、大規模なプログラムにスケールするには難があります。 Zig は新しい言語であるにもかかわらずこの方針を取っていることには疑問があります。

さらに、呼び出されていない関数はコンパイルエラーのチェックも走りません。使用していない変数はコンパイルエラーになるのに、使用していないプライベート関数がエラーにならないというのはよくわからない仕様です。

このような特徴のため、 Zig の特徴の目玉であるコンパイル時変数は、型システムの恩恵をあまり受けられず、ほとんど C のマクロでできるのと同じように感じます。

総評

Rust よりも相当低レベル寄りであると感じました。少し良い感じの機能がちりばめられている印象ですが、成熟しているとはいいがたく、キラーフィーチャーに欠けると思います。

Cじゃだめなんですか？