最近hmatrixで深層学習を実装する機会があったのですが、hmatrixはベクトルと行列しか提供していないので3階以上のテンソルが必要になって困るという場面に出くわしました。そこで自分で長さ付きベクトルを組み合わせてサクッとn階テンソルが作れると便利かな〜と調べていたら（素直にrepaやmassivを使えという話ではあるのですが 
）表現可能関手を使うことでn階テンソルとその演算が楽に（そして抽象的に）実装できるという文献1を見つけたので備忘録も兼ねてまとめておこうと思います。
この記事で紹介したコードは以下のGistで公開しています。
https://gist.github.com/lotz84/78474ac9ee307d50376e025093316d0f
表現可能関手
関手、つまりFunctorのことですが、表現可能関手はその中でも特別な性質を持つものです。この記事は前半と後半に分けて、前半ではHaskellにおける表現可能関手の解説を、後半では表現可能関手を使ったテンソルの実装方法について解説していきたいと思います。
定義と性質
ある型からの関数（つまりReader r）と同型2になるような関手のことを表現可能関手（Representable Functor）34と言います。型クラスとして書くと
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の間の対応5を思い出すと分かりやすいでしょう。実際、
(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)
というような関係が成り立つ6ので（==は同型の意味で使っています）、対数の
\begin{matrix}
\log(ab) &=& \log a + \log b \\
\log a ^ n &=& n\log a
\end{matrix}
という性質と似たような振る舞いをしていることが分かります。
表現可能関手を定義するためにtabulateの代わりにpositionsを採用することも可能です。なぜなら米田の補題7により
{\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を常に成功するように実装することができないからです8。
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もまた表現可能関手になったり9、まだまだ表現可能関手の具体例は尽きませんが、今回はこの辺にして先に進むことにしましょう。
表現可能関手の応用例としてはLogをバケツとして使うことで基数ソートを実装したり10、多相関数a -> bのメモ化を実装したりも9できるのですが、この記事ではテンソルの実装に焦点を当てたいと思います。（最初は書こうかと思ったけど体力がなくなってしまったので興味がある人は誰か記事にしてくれると嬉しいです 
）
テンソルを実装する
さていよいよテンソルを実装していきましょう。以降の実装はAPLicative Programming with Naperian Functorsを参考に進めていきます11。
テンソルは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である場合は値をコピーすることで次元を合わせる
以降では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では単純な値のコピーが発生し効率がよくありませんが、参考元の文献1では別途データ型を定義することでコピーを回避したりもしているので気になる人はぜひ読んでみてください。表現可能関手はReader rに同型であるという強い制約を持つ関手ですが、その分便利な性質をたくさん備えてもいます。他にも表現可能関手を活用した面白い例や応用を見つけたらコメントとかTwitterで教えていただけると嬉しいです。表現可能関手もっと流行れ〜〜！
-
型
A,Bが同型であるとは２つの関数f :: A -> Bとg :: B -> Aでf . g == id,g . f == idとなるようなものが存在することを言います。ざっくり言えばf,gという関数によって情報を失うことなくお互いに変換できるような型同士のことを同型と呼びます。 ↩ -
文献によってはNaperian Functorとも呼ばれることもあります。 ↩
-
数学的な定義はこちら https://en.wikipedia.org/wiki/Representable_functor ↩
-
Representable Functors | Bartosz Milewski's Programming Cafe ↩
-
A Neighborhood of Infinity: Memoizing Polymorphic Functions with High School Algebra and Quantifiers ↩
-
参考元と違うところは、独自の型を使わずに型族を使うことでより構造をわかりやすくしたのと、numpyのbroadcastingと対応させるように
Alignableを拡張したところです。 ↩