テンソルを実装するのに表現可能関手がとても便利な件

最近hmatrixで深層学習を実装する機会があったのですが、hmatrixはベクトルと行列しか提供していないので3階以上のテンソルが必要になって困るという場面に出くわしました。そこで自分で長さ付きベクトルを組み合わせてサクッとn階テンソルが作れると便利かな〜と調べていたら（素直にrepaやmassivを使えという話ではあるのですが ）表現可能関手を使うことでn階テンソルとその演算が楽に（そして抽象的に）実装できるという文献¹を見つけたので備忘録も兼ねてまとめておこうと思います。

この記事で紹介したコードは以下のGistで公開しています。
https://gist.github.com/lotz84/78474ac9ee307d50376e025093316d0f

表現可能関手

関手、つまりFunctorのことですが、表現可能関手はその中でも特別な性質を持つものです。この記事は前半と後半に分けて、前半ではHaskellにおける表現可能関手の解説を、後半では表現可能関手を使ったテンソルの実装方法について解説していきたいと思います。

定義と性質

ある型からの関数（つまりReader r）と同型²になるような関手のことを表現可能関手（Representable Functor）³⁴と言います。型クラスとして書くと

class Functor f => Representable f where
  type Log f
  index   :: f a -> (Log f -> a)
  tabulate :: (Log f -> a) -> f a

  positions :: f (Log f)
  tabulate h = fmap h positions
  positions  = tabulate id

このようになり、indexとtabulateが関手と関数の同型対応を与えています。

関数の引数となる型にLogという名前がついている気持ちは$a^b$とb -> aの間の対応⁵を思い出すと分かりやすいでしょう。実際、

(f a, g a) == (Log f -> a, Log g -> a) == Either (Log f) (Log g) -> a
=> Log (f a, g a) == Either (Log f) (Log g)

n -> f a == n -> Log f -> a == (n, Log f) -> a
=> Log (n -> f a) == (n, Log f)

というような関係が成り立つ⁶ので（==は同型の意味で使っています）、対数の

\begin{matrix}
\log(ab) &=& \log a + \log b \\
\log a ^ n &=& n\log a
\end{matrix}

という性質と似たような振る舞いをしていることが分かります。

表現可能関手を定義するためにtabulateの代わりにpositionsを採用することも可能です。なぜなら米田の補題⁷により

{\rm Nat}({\rm Hom}(A, -), F) \cong F A

という同型対応が存在し、この対応がまさにtabulateとpositionsが相互に変換可能であることを表しています。実際Representableのインスタンスを定義する時はpositionsを実装する方が楽なことの方が多いので覚えておくと良いでしょう。

表現可能関手はReader rと同型になるのでそこから色々な性質を受け継ぐことが出来ます。例えば表現可能関手であればApplicativeのインスタンスを定義することは簡単です。

instance Representable f => Applicative f where
  pure a = tabulate $ \_ -> a
  f <*> g = tabulate $ \i -> index f i $ index g i

ただこのやり方だと UndecidableInstances が必要になってしまうので、実際は個別に定義するのが良いでしょう。もちろんモナドを定義することも可能です。

さらに表現可能関手であればTraversableの双対であるDistributiveにもなります。Distributiveは

distribute :: (Functor f, Distributive g) => f (g a) -> g (f a)

というメソッドを備えた型クラスで、表現可能関手であれば

distribute fga = tabulate $ \i -> fmap (`index` i) fga

として実装することができます。さらに表現可能関手同士のdistributeであれば

distribute = tabulate . fmap tabulate . flip . fmap index . index

と実装することができるので、実質flipと同じ操作であることがわかります。後で見るように、このインデックスの順番を取り替えるだけという操作が直接そのまま行列の転置の実装に対応します。

表現可能関手のインスタンス

それでは具体例を見ていきましょう。

(->) r

表現可能関手はReader rと同型になるような関手だったので関数の型は自然に表現可能関手となります。

instance Representable ((->) r) where
  type Log ((->) r) = r
  index = ($)
  positions = id

Identity

Identityは最も基本的な関手ですが、これも表現可能関手と捉えることが出来ます。

instance Representable Identity where
  type Log Identity = ()
  index a _ = runIdentity a
  positions = Identity ()

同様に数値型にモノイドの構造を与えるSumやProduct、幽霊型を付け加えるTaggedなども表現可能関手とみなすことが出来ます。

対角関手

対角関手を以下のように定義します。

data Diag a = Diag a a

instance Functor Diag where
  fmap f (Diag a b) = Diag (f a) (f b)

すると以下のようにこれは表現可能関手のインスタンスとなります。

instance Representable Diag where
  type Log Diag = Bool
  index (Diag a b) c = if c then b else a
  positions = Diag False True

同様にComplexも表現可能関手のインスタンスにすることが可能です。LogとしてBoolが使われているのは単に"２個の値を区別できる型"なら何でもよくて真偽値である必要はありません。これを一般化して"n個の値を区別できる型"を使うことでVector nを表現可能関手のインスタンスにすることが出来ます。

Vector n a

"n個の値を区別できる型"としてFinite nを使います。これはn未満の自然数を表す型です。

import qualified Data.Vector.Sized as V

instance KnownNat n => Representable (V.Vector n) where
  type Log (V.Vector n) = Finite n
  index = V.index
  positions = V.generate id

これでベクトルを表現可能関手のインスタンスにすることが出来ました。この調子で2階以上のテンソルのインスタンスも作っていきたいところですが、そのためには表現可能関手の合成を表現可能関手とみなす事ができれば簡単です。なので焦らずにもう少しいくつかの例を見ていくことにしましょう。

Stream a

n個の要素を持つ型を表現可能関手に出来たので無限個の要素を持つ型を表現可能関手のインスタンスにすることが出来ないのかと考えるのは自然でしょう。まずは無限個の要素を持つデータ型Streamを定義してみましょう。

data Stream a = Cons a (Stream a)

instance Functor Stream where
  fmap f (Cons a as) = Cons (f a) (fmap f as)

sIndex :: Stream a -> Natural -> a
sIndex (Cons a _ ) 0 = a
sIndex (Cons _ as) n = sIndex as (n-1)

ただしNaturalは自然数を表す型です。
このStream aを表現可能関手のインスタンスにすることが出来ます。

instance Representable Stream where
  type Log Stream = Natural
  index = sIndex
  positions = fix $ \xs -> Cons 0 (fmap (+1) xs)

同じ要領で通常のリスト[a]も表現可能関手のインスタンスに出来そうな気もするかもしれませんが実は出来ません。なぜならリストは長さが有限ですがそれがどれぐらいの大きさになるかを事前に型から知ることが出来ないのでindexを常に成功するように実装することができないからです⁸。

Product

表現可能関手の積Productはまた表現可能関手になります。

instance (Representable f, Representable g) => Representable (Product f g) where
  type Log (Product f g) = Either (Log f) (Log g)
  index (Pair fa _) (Left  i) = index fa i
  index (Pair _ ga) (Right i) = index ga i
  positions = Pair (tabulate Left) (tabulate Right)

積のLogを取ると和Eitherが出てくるのは面白いですよね。

Compose

表現可能関手の合成Composeもまた表現可能関手になります。

instance (Representable f, Representable g) => Representable (Compose f g) where
  type Log (Compose f g) = (Log f, Log g)
  index fga (i, j) = (`index` j) . (`index` i) $ getCompose fga
  tabulate = Compose . tabulate . fmap tabulate . curry

これを使って Compose (Vector m) (Vector n) のようにベクトルを合成していけば高階テンソルの型を作ることができ、表現可能関手としての性質も利用することが出来るはずです。

他にも木構造やさらには表現可能関手のCofreeもまた表現可能関手になったり⁹、まだまだ表現可能関手の具体例は尽きませんが、今回はこの辺にして先に進むことにしましょう。

表現可能関手の応用例としてはLogをバケツとして使うことで基数ソートを実装したり¹⁰、多相関数a -> bのメモ化を実装したりも⁹できるのですが、この記事ではテンソルの実装に焦点を当てたいと思います。（最初は書こうかと思ったけど体力がなくなってしまったので興味がある人は誰か記事にしてくれると嬉しいです ）

テンソルを実装する

さていよいよテンソルを実装していきましょう。以降の実装はAPLicative Programming with Naperian Functorsを参考に進めていきます¹¹。

テンソルはVector nの合成として実装します。しかし単純に合成で書いてしまうと型の演算と相性が悪くなってしまうので、テンソルの各次元を型レベル自然数のリストで表し、型族によってVector nの合成に変換するという方針で実装してみましょう。加えてTypeFamilyDependencies拡張を使って型レベル自然数からの単射であることも明示しておきます。

type family Tensor (xs :: [Nat]) = r | r -> xs where
  Tensor '[] = Identity
  Tensor (n ': ns) = Compose (Tensor ns) (V.Vector n)

例えばTensor '[n] aはn次元のベクトルで、Tensor '[n, m] aはm✕nの行列に対応します。ここでmとnの順番がひっくり返ってることに注意してください。実はこっちの順番のほうが後にMaxやAlignableといった型レベルの演算を定義するときに実装が楽になるのです。

type Vector n a = Tensor '[n] a
type Matrix m n a = Tensor '[n, m] a

まずは具体的なテンソルの値を作れるようにリストから変換する関数を作ってみましょう。

class FromList f where
  fromList :: [a] -> f a

instance FromList Identity where
  fromList [a] = Identity a

instance forall n t. (KnownNat n, FromList t) => FromList (Compose t (V.Vector n)) where
  fromList as =
    let n = fromIntegral $ natVal (Proxy @n)
     in Compose . fromList . map (fromJust . V.fromList) $ chunksOf n as

fromListはリストからテンソルに変換する関数です。テンソルの定義は基底部と再起部のように別れていますが、型によってこれらの場合分けを行うために型クラスを利用しています。

さっそくベクトルの内積を定義してみましょう。

liftR2 :: Representable f => (a -> b -> c) -> f a -> f b -> f c
liftR2 f fa fb = tabulate $ \i -> f (index fa i) (index fb i)

dot :: (KnownNat n, Num a) => Vector n a -> Vector n a -> a
dot = (sum.) . liftR2 (*)

liftR2は表現可能関手に対して一般的に定義できる関数でzipWithのような役割を果たします。ベクトルの内積dotはliftR2で要素同士の積をとりFoldableの性質を利用して足し合わせているだけです。

試しに実行してみると、

> xs = fromList [1,2,3] :: Vector 3 Int
> xs `dot` xs
14

となりうまく実装できていそうです。

次に行列の転置を定義してみましょう。

transpose :: (KnownNat m, KnownNat n) => Matrix m n a -> Matrix n m a
transpose (Compose (Compose xs)) = Compose (Compose (fmap distribute xs))

Composeをうまく剥がしてVector m (Vector n a)の型に変換した後にdistributeを適用しているだけです。fmapはIdentityをまたぐために使っています。

試しに実行してみると、

> mat = fromList [1..6] :: Matrix 3 2 Int
> mat
Compose (Compose (Identity [[1,2],[3,4],[5,6]]))

> transpose mat
Compose (Compose (Identity [[1,3,5],[2,4,6]]))

となりうまく転置出来ていることが分かります。ちなみにdistributeと同じように関手の順番を入れ替えるTraversableのsequenceを二重リストに適用すると

> sequence [[1, 2], [3, 4], [5, 6]]
[[1,3,5],[1,3,6],[1,4,5],[1,4,6],[2,3,5],[2,3,6],[2,4,5],[2,4,6]]

のように多重ループと同じふるまいとなってしまい転置とはかけ離れてしまう事を考えると、distributeが転置といかにうまく対応しているかが分かると思います。

これらを利用して行列積を定義してみましょう。

matmul :: (KnownNat p, KnownNat q, KnownNat r, Num a)
       => Matrix p q a -> Matrix q r a -> Matrix p r a
matmul (Compose a) b' =
  let (Compose b) = transpose b'
      vec = Compose . Identity
   in tabulate $ \(((), j), i) -> (vec $ index a ((), j)) `dot` (vec $ index b ((), i))

ComposeとIdentityを付けたり剥がしたりするところが少し煩雑ですが、大事な部分はtabulateを使うことで(i, j)番目の要素にだけ注目して実装できているというところです。誤解を恐れずに言えば表現可能関手を使えば関手の形のデータを直接利用することが出来るようになるので、各パーツにおける処理さえ記述すれば、パーツに分解して処理した後に全体を組み立てるという操作は表現可能関手としての性質に任せてしまうことが出来るのです。便利ですね、表現可能関手！

ランク多相、そしてBroadcasting

最後にランク多相な演算子を実装してみましょう。ランク多相とはテンソルのランク（テンソルの階数）に対して多相になっていることを指し、ランクN多相とは別の概念です。例えば平方数を計算するsquareという関数は、スカラーに適用すると

square 2 = 4

となり、行列に適用すると

square [[1, 2]  = [[1,  4]
       ,[3, 4]]   ,[9, 16]]

のように成分ごとに適用されるような振る舞いが期待されるかと思います。これがランク多相という概念です。

一引数の関数であれば、

unary :: Functor (Tensor ns) => (a -> b) -> Tensor ns a -> Tensor ns b
unary = fmap

のようにfmapとして簡単に実現することが出来ます。

２つ引数を取る場合はどうでしょう？例えばベクトル同士の足し算は

[1, 2] + [3, 4] = [3, 6]

というように成分ごとの足し算が期待されると思います。加えて

[1, 2] + 1 = [2, 3]

このように次元が違う場合でも演算を許容すると、実際にプログラミングする際には非常に便利です。

このような概念はNumPyではBroadcastingという名前で使われており、以下のルールで行われています。

1. ランクが異なる場合は1次元ずつ加えることでランクが低い方を高い方に合わせる
2. 対応する次元が同じであれば良し。片方が1である場合は値をコピーすることで次元を合わせる

以降ではBroadcastingを実現したランク多相な二項演算子を実装していきたいと思います。

まず、同じ値でテンソルを埋める関数replicateTを実装します。

class Shapely ns where
  replicateT :: a -> Tensor ns a

instance Shapely '[] where
  replicateT a = Identity a

instance (KnownNat n, Shapely ns, Representable (Tensor ns)) => Shapely (n ': ns) where
  replicateT a = Compose (replicateT (tabulate (const a)))

これを使えば、

> replicateT 0 :: Matrix 2 3 Int
Compose (Compose (Identity [[0,0,0],[0,0,0]]))

のようにゼロ埋めした行列などを簡単に作れるようになります。

次にBroadcastingした結果のランクを計算するMaxという型族を実装します。

type family Max xs ys where
  Max '[] '[] = '[]
  Max (x:xs) '[] = x ': Max xs '[]
  Max '[] (y:ys) = y ': Max ys '[]
  Max (x:xs) (y:ys) = (If (x <=? y) y x) ': Max xs ys

最後にMaxで計算したランクに変形するためのalignという関数を定義します。

class Alignable (ns :: [Nat]) (ms :: [Nat]) where
  align :: Tensor ns a -> Tensor ms a

instance Alignable '[] '[] where
  align = id

instance Alignable xs ys => Alignable (x ': xs) (x ': ys) where
  align (Compose t) = Compose (align t)

instance (KnownNat y, Functor (Tensor xs), Alignable xs ys) => Alignable (1 ': xs) (y ': ys) where
  align =
    let y = natVal (Proxy @y)
     in Compose . align . fmap (V.replicate . V.head) . getCompose

instance (KnownNat n, Shapely ns, Representable (Tensor ns)) => Alignable '[] (n ': ns) where
  align = replicateT . runIdentity

alignを実装するためのAlignableという型クラスは4つの部分からなります。まずは基底部に対応する部分で、これは恒等変換になります。次に対応する次元が同じ値を持つときですが、これは何もせずに次の次元にalignを適用していきます。次に変換元の次元が1の場合、これは変換先の次元分だけ元の値をコピーします。最後に変換元のランクが低く次元が先になくなってしまった場合、最後のスカラー値を残りの次元分コピーします。

以上を組み合わせてBroadcastingを実装したランク多相な二項演算子の完成です。

binary :: (Max xs ys ~ zs, Alignable xs zs, Alignable ys zs, Representable (Tensor zs))
       => (a -> b -> c)
       -> Tensor xs a -> Tensor ys b -> Tensor zs c
binary f xs ys = liftR2 f (align xs) (align ys)

早速試してみましょう。

> binary (+) (fromList [1..4] :: Matrix 2 2 Int) (fromList [1..4] :: Matrix 2 2 Int)
Compose (Compose (Identity [[2,4],[6,8]]))

> binary (+) (fromList [1, 2] :: Vector 2 Int) (Identity 1)
Compose (Identity [2,3])

> binary (*) (fromList [1, 2] :: Vector 2 Int) (fromList [1..8] :: Tensor '[2, 2, 2] Int)
Compose (Compose (Compose (Identity [[[1,4],[3,8]],[[5,12],[7,16]]])))

うまく行ってそうですね！

最後に最近話題になったBroadcastingの罠っぽい挙動も再現してみましょう。

Here's an amazing fact about Numpy. Don't fall into this trap! pic.twitter.com/7I64zLEfSN

— François Chollet (@fchollet) September 15, 2019

> x = replicateT 1 :: Tensor '[1, 32] Int
> y = replicateT 1 :: Tensor '[32] Int
> z = binary (+) x y
> sum z
2048

無事（？）再現できましたね

まとめ

表現可能関手を活用して（最後は型レベルプログラミングでしたがｗ）テンソルを簡単に実装することが出来ました。Broadcastingでは単純な値のコピーが発生し効率がよくありませんが、参考元の文献¹では別途データ型を定義することでコピーを回避したりもしているので気になる人はぜひ読んでみてください。表現可能関手はReader rに同型であるという強い制約を持つ関手ですが、その分便利な性質をたくさん備えてもいます。他にも表現可能関手を活用した面白い例や応用を見つけたらコメントとかTwitterで教えていただけると嬉しいです。表現可能関手もっと流行れ〜〜！

1. APLicative Programming with Naperian Functors (PDF) ↩ ↩²

2. 型A, Bが同型であるとは２つの関数f :: A -> Bとg :: B -> Aでf . g == id, g . f == id となるようなものが存在することを言います。ざっくり言えばf, gという関数によって情報を失うことなくお互いに変換できるような型同士のことを同型と呼びます。 ↩

3. 文献によってはNaperian Functorとも呼ばれることもあります。 ↩

4. 数学的な定義はこちら https://en.wikipedia.org/wiki/Representable_functor ↩

5. 代数的データ型と初等代数学 - ryota-ka's blog ↩

6. The Comonad.Reader » Representing Applicatives ↩

7. 米田の補題 | 壱大整域 ↩

8. Representable Functors | Bartosz Milewski's Programming Cafe ↩

9. A Neighborhood of Infinity: Memoizing Polymorphic Functions with High School Algebra and Quantifiers ↩ ↩²

10. https://twitter.com/lotz84_/status/1347789776309551104 ↩

11. 参考元と違うところは、独自の型を使わずに型族を使うことでより構造をわかりやすくしたのと、numpyのbroadcastingと対応させるようにAlignableを拡張したところです。 ↩