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

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

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

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

定義と性質
--------

ある型からの関数（つまり`Reader r`）と同型[^3]になるような関手のことを**表現可能関手（Representable Functor）**[^1][^2]と言います。型クラスとして書くと

```hs
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`の間の対応[^4]を思い出すと分かりやすいでしょう。実際、

```hs

(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)
```

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

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

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

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

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

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

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

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

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

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

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

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

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

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

```hs
distribute = tabulate . fmap tabulate . flip . fmap index . index
```

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

表現可能関手のインスタンス
---------------------
それでは具体例を見ていきましょう。

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

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

### `Identity`
[`Identity`](http://hackage.haskell.org/package/base-4.11.0.0/docs/Data-Functor-Identity.html)は最も基本的な関手ですが、これも表現可能関手と捉えることが出来ます。

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

同様に数値型にモノイドの構造を与える[`Sum`や`Product`](http://hackage.haskell.org/package/base-4.12.0.0/docs/Data-Monoid.html#g:3)、幽霊型を付け加える[`Tagged`](http://hackage.haskell.org/package/tagged-0.8.6/docs/Data-Tagged.html)なども表現可能関手とみなすことが出来ます。

### 対角関手

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

```hs
data Diag a = Diag a a

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

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

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

同様に[`Complex`](http://hackage.haskell.org/package/base-4.12.0.0/docs/Data-Complex.html)も表現可能関手のインスタンスにすることが可能です。`Log`として`Bool`が使われているのは単に"２個の値を区別できる型"なら何でもよくて真偽値である必要はありません。これを一般化して"n個の値を区別できる型"を使うことで[`Vector n`](http://hackage.haskell.org/package/vector-sized)を表現可能関手のインスタンスにすることが出来ます。

### `Vector n a`
"n個の値を区別できる型"として[`Finite n`](https://hackage.haskell.org/package/finite-typelits-0.1.4.2/docs/Data-Finite.html)を使います。これはn未満の自然数を表す型です。

```hs
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`を定義してみましょう。

```hs
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`](http://hackage.haskell.org/package/base-4.12.0.0/docs/Numeric-Natural.html)は自然数を表す型です。
この`Stream a`を表現可能関手のインスタンスにすることが出来ます。

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

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

### `Product`
表現可能関手の積[`Product`](http://hackage.haskell.org/package/base-4.12.0.0/docs/Data-Functor-Product.html)はまた表現可能関手になります。

```hs
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`](http://hackage.haskell.org/package/base-4.12.0.0/docs/Data-Functor-Compose.html)もまた表現可能関手になります。

```hs
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)` のようにベクトルを合成していけば高階テンソルの型を作ることができ、表現可能関手としての性質も利用することが出来るはずです。

:earth_asia: 

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

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

テンソルを実装する
==============
さていよいよテンソルを実装していきましょう。以降の実装は[APLicative Programming with Naperian Functors](https://www.cs.ox.ac.uk/people/jeremy.gibbons/publications/aplicative.pdf)を参考に進めていきます[^11]。

テンソルは`Vector n`の合成として実装します。しかし単純に合成で書いてしまうと型の演算と相性が悪くなってしまうので、テンソルの各次元を型レベル自然数のリストで表し、型族によって`Vector n`の合成に変換するという方針で実装してみましょう。加えて[TypeFamilyDependencies拡張](https://ghc.gitlab.haskell.org/ghc/doc/users_guide/glasgow_exts.html#extension-TypeFamilyDependencies)を使って型レベル自然数からの単射であることも明示しておきます。

```hs
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`といった型レベルの演算を定義するときに実装が楽になるのです。

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

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

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

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

```hs
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`](http://hackage.haskell.org/package/adjunctions-4.4/docs/Data-Functor-Rep.html#v:liftR2)は表現可能関手に対して一般的に定義できる関数で`zipWith`のような役割を果たします。ベクトルの内積`dot`は`liftR2`で要素同士の積をとり`Foldable`の性質を利用して足し合わせているだけです。

試しに実行してみると、

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

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

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

```hs
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`をまたぐために使っています。

試しに実行してみると、

```hs
> 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`を二重リストに適用すると

```hs
> 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`が転置といかにうまく対応しているかが分かると思います。

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

```hs
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多相](https://ghc.gitlab.haskell.org/ghc/doc/users_guide/glasgow_exts.html#extension-RankNTypes)とは別の概念です。例えば平方数を計算する`square`という関数は、スカラーに適用すると

```hs
square 2 = 4
```

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

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

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

一引数の関数であれば、

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

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

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

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

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

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

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

このような概念はNumPyでは[Broadcasting](https://docs.scipy.org/doc/numpy/user/basics.broadcasting.html)という名前で使われており、以下のルールで行われています。

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

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

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

```hs
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)))
```

これを使えば、

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

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

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

```hs
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`という関数を定義します。

```hs
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を実装したランク多相な二項演算子の完成です。

```hs
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)
```

早速試してみましょう。

```hs
> 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の罠っぽい挙動も再現してみましょう。

<blockquote class="twitter-tweet"><p lang="en" dir="ltr">Here&#39;s an amazing fact about Numpy. Don&#39;t fall into this trap! <a href="https://t.co/7I64zLEfSN">pic.twitter.com/7I64zLEfSN</a></p>&mdash; François Chollet (@fchollet) <a href="https://twitter.com/fchollet/status/1173306347175075841?ref_src=twsrc%5Etfw">September 15, 2019</a></blockquote> <script async src="https://platform.twitter.com/widgets.js" charset="utf-8"></script>

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

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

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

[^1]: 文献によってはNaperian Functorとも呼ばれることもあります。
[^2]: 数学的な定義はこちら https://en.wikipedia.org/wiki/Representable_functor
[^3]: 型`A`, `B`が同型であるとは２つの関数`f :: A -> B`と`g :: B -> A`で`f . g == id`, `g . f == id` となるようなものが存在することを言います。ざっくり言えば`f`, `g`という関数によって情報を失うことなくお互いに変換できるような型同士のことを同型と呼びます。
[^4]: [代数的データ型と初等代数学 - ryota-ka's blog](https://ryota-ka.hatenablog.com/entry/2018/07/09/110000)
[^5]: [The Comonad.Reader » Representing Applicatives](http://comonad.com/reader/2013/representing-applicatives/)
[^6]: [米田の補題 | 壱大整域](http://alg-d.com/math/kan_extension/yoneda.pdf)
[^7]: [APLicative Programming with Naperian Functors](https://www.cs.ox.ac.uk/people/jeremy.gibbons/publications/aplicative.pdf) (PDF)
[^8]: [Representable Functors | Bartosz Milewski's Programming Cafe](https://bartoszmilewski.com/2015/07/29/representable-functors/)
[^9]: <https://twitter.com/lotz84_/status/1347789776309551104>
[^10]: [A Neighborhood of Infinity: Memoizing Polymorphic Functions with High School Algebra and Quantifiers](http://blog.sigfpe.com/2009/11/memoizing-polymorphic-functions-with.html)
[^11]: 参考元と違うところは、独自の型を使わずに型族を使うことでより構造をわかりやすくしたのと、numpyのbroadcastingと対応させるように`Alignable`を拡張したところです。