Zenn
😽

C++で(One-shot) Algebraic Effectsで自動微分をする

2023/04/16に公開
C++
tech

はじめに

Algebraic Effects (and Handlers)は大雑把に言うと「もとの場所に戻れる例外」で、戻らなくても良いし、また戻った後に何か動作させることもできるし、（multi-shotであれば）複数回もとの場所に戻ることもできます。

例の如くDan Abramov氏の記事
https://overreacted.io/ja/algebraic-effects-for-the-rest-of-us/
の記法を使うと

// （1回）戻る場合
try {
  makeFriends(arya, gendry);
} handle (effect) {
  if (effect === 'ask_name') {
    resume with 'Arya Stark';
  }
}

// 戻らない、つまり普通の例外と同じ
try {
  makeFriends(arya, gendry);
} handle (effect) {
  if (effect === 'ask_name') {
    console.error('ERROR');
  }
}

// 戻った後ロギング
try {
  makeFriends(arya, gendry);
} handle (effect) {
  if (effect === 'ask_name') {
    resume with 'Arya Stark';
    console.log('resumed');
  }
}

と言った感じです（JSには存在しない機能なので全てフィーリングですが）。

この機能を実装している言語としては、Eff、Koka、Effekt、そしてUnisonなどがあり、最近ではOCamlが対応して話題になりました。

実際どんな風に使ってどんなことができるの？というのはEffektのチュートリアルや幾つかの論文[1]が参考になります。非同期プログラミング、パーサー、自動微分、確率的プログラミングなどの応用があります。

そんなAlgebraic Effectsはライブラリレベルでの実装もいくつかあり、Lua、Ruby、Rustなどで実装されています。

今回はC++でAlgebraic Effectsを使えるライブラリ、cpp-effectsを使ってみます。

cpp-effects

https://github.com/maciejpirog/cpp-effects
cpp-effectsはこちらの論文で提案された、C++でAlgebraic Effectsを実現するライブラリです。
C/C++では、libhandlerがありますが、こちらはよりC++らしいHigh LevelなAPIを提供しています。
examplesディレクトリにいくつか例がありますが、ここでは（そこにはない）自動微分を実装してみます。

自動微分

ここでの実装はEffekt言語のチュートリアルを基にしています。
https://effekt-lang.org/docs/casestudies/ad
また自動微分そのものについては説明はしません。

#include <memory>
#include <iostream>
#include "cpp-effects/cpp-effects.h"
#include "cpp-effects/clause-modifiers.h"

namespace eff = cpp_effects;

次に「変数」的なものを定義します。

struct Num {
    double value;
    std::shared_ptr<double> d;
};

これを発生させるEffectは以下のように定義します。

struct NumOp : eff::command<Num> {
    double x;
};

auto num(double x) { 
    return eff::invoke_command(NumOp{{}, x});
}

また足し算と掛け算のEffectも定義します。演算子オーバーロードもしておきます。

struct AddOp : eff::command<Num> {
    Num x;
    Num y;
};

auto add(const Num& x, const Num& y) { 
    return eff::invoke_command(AddOp{{}, x, y});
}

struct MulOp : eff::command<Num> {
    Num x;
    Num y;
};

auto mul(const Num& x, const Num& y) {
    return eff::invoke_command(MulOp{{}, x, y});
}

auto operator+(const Num& x, const Num& y) {
    return add(x, y);
}

auto operator*(const Num& x, const Num& y) {
    return mul(x, y);
}

これを用いて以下の式を考えます。

auto prog(const Num& x) { 
    return (num(3.0) * x) + (x * x * x);
}

順伝播は以下とします。

auto forward() {
    return prog(Num{ 2.0, std::make_shared<double>(1.0) });
}

これを計算すると3 * 2 + 2 * 2 * 2で14になってほしいわけです。
そこで順伝播のhandlerを定義します。

class FP : public eff::flat_handler<Num, eff::plain<NumOp>, eff::plain<AddOp>,  eff::plain<MulOp>> {
    private:
        Num handle_command(NumOp x) override {
            return Num{x.x, std::make_shared<double>(1.0)};
        }
        Num handle_command(AddOp x) override {
            return Num{x.x.value + x.y.value, std::make_shared<double>(*(x.x.d) + *(x.y.d))};
        }
        Num handle_command(MulOp x) override {
            return Num{x.x.value * x.y.value, std::make_shared<double>(*(x.x.d) * x.y.value + *(x.y.d) * x.x.value)};
        }
};

このように各Effectのhandlerを書きます[2]。
使ってみましょう。

std::cout << eff::handle<FP>(forward).value << std::endl;
// output: 14

次に逆伝播を考えます。式を微分すると3.0 + 3x^2になるはずです。x = 2.0を代入すれば15になるはずです。
handlerを定義します。

class BP : public eff::handler<void, Num, NumOp, AddOp, MulOp> {
    private:
        void handle_command(NumOp x, eff::resumption<void(Num)> r) override {
            std::move(r).tail_resume(Num{x.x, std::make_shared<double>(0.0)});
        }
        void handle_command(AddOp x, eff::resumption<void(Num)> r) override {
            auto z = Num{x.x.value + x.y.value, std::make_shared<double>(0)};
            std::move(r).resume(z);
            push(x.x, *(z.d));
            push(x.y, *(z.d));
        }
        void handle_command(MulOp x, eff::resumption<void(Num)> r) override {
            auto z = Num{x.x.value * x.y.value, std::make_shared<double>(0)};
            std::move(r).resume(z);
            push(x.x, *(z.d) * x.y.value);
            push(x.y, *(z.d) * x.x.value);
        }
        void handle_return(Num) override {
        }
};

今回はresume後にgrad部分の更新が必要なので、flat_handlerではなくhandlerを使用しています。handle_returnはvalue handler/return clause[3]です。ちなみにflat_handlerを使う場合のvalue handlerはidentityです。
こうしてhandlerを入れ替えて

auto backward() {
    auto input = Num{ 2.0, std::make_shared<double>(0.0) };
    eff::handle<BP>(
        [&input](){
            return [&input]() { 
                auto a = prog(input);
                push(a, 1.0);
                return a;
            }();
        }
    );
    return *(input.d);
}
std::cout << backward() << std::endl;
// 15

と逆伝播が計算できました。

おわりに

このライブラリはMulti-shotな継続には対応していません。対応されるとパーサーが書けるのでより楽しめますね。
とはいえOne-shotでも多くのことに使えますし、Algebraic EffectsはC++を含む多くの言語でライブラリがあるので、面白いですね（？）

脚注

  1. Algebraic Effects for Functional ProgrammingやProgramming with Algebraic Effects and Handlers ↩︎

  2. Effectと言うかOperationと言うかライブラリ的にはcommandとそのcommand clause ↩︎

  3. 自信はありませんがvalue handler/return clauseは同じもののはず ↩︎

Discussion