More than 1 year has passed since last update.

Numo::NArray で行列の積（あるいはテンソルの縮約）を計算

Last updated at 2022-01-09Posted at 2021-11-20

Ruby で高速に多次元数値配列を扱いたいときに Numo::NArray を使っているのだが、この配列は基本的に「要素毎」の演算となっている。では NArray で行列の積などを計算したいときはどうすればいいのか気になったので、調べて整理した。

TL;DR

よく使う積はメソッドが用意されていて、行列の積なら a.dot(b) でいい
自分でやる場合、行列の積 $C_{ij} = \sum_{k} A_{ik} B_{kj}$ は c = a[false, :new].mulsum(b, -2) で計算できる（ a, b は2次元配列）
より一般にテンソルの縮約を計算する場合、掛け合わせるテンソルの添字が揃うように配列の次元を追加し、 #mulsum で縮約する次元を指定して積和計算する

（前提知識）ビューの次元指定について

a[false, :new, true] と書かれていても意味がわからないときのためのメモ。

クリックして詳細を表示

NArray は多次元配列を扱うので、例えば3次元配列 a = Numo::Int64.new(2, 4, 6).indgen に対して a[1, 2..3, 0..-1] のように次元毎にインデックスや範囲を指定して部分配列を切り出せる¹
- 部分配列は「ビュー」（view）と呼び、これを通して元の配列の対応部を読み書きできる
- 次元毎の場合、指定する引数の個数は配列の次元数 a.ndim と合っていないといけない（例なら3つ）
  …というのはあくまで基本ルールで、 false や :new は別途考える必要がある
引数のうち 0..-1 など全範囲を表すものは true と略記できる
0個以上の連続する true は、1ヶ所だけ false で置き換えられる（rest dimensions）
- 任意次元の配列を扱いたいときに特に便利
:new を追加すると、そこに大きさ1の次元を追加できる
- 追加なので「引数の個数を配列の次元数と合わせる」というルールではカウントされず、何ヶ所でもいい
- 同様のことは #expand_dims というメソッドでも可能（ただし1回につき1ヶ所のみ）

従って以下は全て同じビューを得られる。（一部の方法は a.ndim == 3 を前提としている）

# a = Numo::Int64.new(2, 4, 6).indgen

a[false, :new, true]  #=> Numo::Int64(view)#shape=[2,4,1,6] [...]
a[false, true].expand_dims(-2)
a[true, true, true].expand_dims(-2)
a[0..-1, 0..-1, 0..-1].expand_dims(-2)
a.view.expand_dims(-2)
a.expand_dims(-2)

なお、特別な意味を持つ引数は他にもあるみたいだが（該当コード）、ドキュメントなどでの使用例が少ないので、本記事では true, false, :new で統一する。

機能	引数に与えるオブジェクト
全範囲	`true`, `nil`, `:all`, `:*`
全範囲・逆順	`:reverse`
全範囲・集計フラグ設定	`:reduce`, `:sum`, `:+`
残余次元	`false`, `:rest`, `:~`
次元追加	`:new`, `:-`

1次元配列同士の積

手始めに、ベクトルに見立てた2つの1次元配列の演算を扱う。

require 'numo/narray'

u = Numo::Int64[1, 2, 3, 4]
v = Numo::Int64[1, 10, 100, 1000]

単純に配列同士を掛け算すると、要素毎の積（アダマール積）が返る。

u * v
#=>
# Numo::Int64#shape=[4]
# [1, 20, 300, 4000]

内積（ドット積）を計算するには、要素毎の積の総和をとればいい： $a = \sum_i u_i v_i$ 。これは前述の計算の後に #sum を作用させてもいいし、 #mulsum で一気に計算することもできる。

(u * v).sum  #=> 4321
u.mulsum(v)  #=> 4321

もうひとつ別の積として外積（二項積、直積、テンソル積）²も考えられる： $A_{ij} = u_i v_j$ 。これは九九の表を作るイメージでよく、ベクトルの大きさは異なっていてもいい。 NArray で実現するには、大きさ1の次元を追加して添字が被らないようにして積をとる³。

                      #    i, j
u[false, :new].shape  #=> [4, 1]
v             .shape  #=>    [4]

u[false, :new] * v
#=> 
# Numo::Int64#shape=[4,4]
# [[1, 10, 100, 1000], 
#  [2, 20, 200, 2000], 
#  [3, 30, 300, 3000], 
#  [4, 40, 400, 4000]]

行列の積の計算では、この内積と外積の考え方を組み合わせる。

2次元配列による行列の積

行列の積は $C_{ij} = \sum_{k} A_{ik} B_{kj}$ と計算する。図で説明しているもの⁴を見るとわかりやすいが、 C の各要素は A の行と B の列との内積となっている（また、これを全ての行と列の組に対して計算しているのは外積と似ている）。内積は前節で扱ったように積の総和という形に分解して考えていい。

NArray は基本的に要素毎の積を計算するので、行と列というように次元がずれているものは計算できない。そこで配列に大きさ1の次元を追加して、内積をとりたい次元を揃えることで解決する。行列の積であれば、上式 A の末尾に次元をひとつ追加すると k は後ろから2次元目で揃う（添字が早く回るのは後ろ側なので、次元は後ろから数えて合わせること）。あとは内積計算を完了するために k の次元で総和をとればいい。
ちなみに k の次元を揃えたことで i, j の次元は相手側の大きさが1になっている。これは前節の外積と同じで、全ての (i, j) の組について計算することを表す。

require 'numo/narray'

a = Numo::Int64[[1, 2], [3, 4]]
b = Numo::Int64[[1, 10], [100, 1000]]

                      #    i, k, j
a[false, :new].shape  #=> [2, 2, 1]
b             .shape  #=>    [2, 2]

a[false, :new] * b  # kの次元を揃えた状態で要素毎の積を計算する
#=>
# Numo::Int64#shape=[2,2,2]
# [[[1, 10],
#   [200, 2000]],
#  [[3, 30],
#   [400, 4000]]]

(a[false, :new] * b).sum(-2)  # 後ろから2番目(k)の次元で和をとる
#=>
# Numo::Int64#shape=[2,2]
# [[201, 2010],
#  [403, 4030]]

a[false, :new].mulsum(b, -2)
#=> (same)

積と和を分けて計算すると、積をとった段階で3次元配列が作られメモリを多く消費する。そのため #mulsum で一気に計算したほうが効率がいい。

一般化

数式が添字表記で、どの次元で和をとるのかわかっていれば、同様に次元追加（および必要なら転置）して #mulsum で計算すればいい。物理学で式をテンソルで表す場合などによく登場し、総和記号を省略することもある。

例えば $b_{pq} = \sum_{i,j} r_{pi} r_{qj} a_{ij}$ （2階テンソルの座標変換をイメージした式）なら、要素を表す添字は $pq, pi, qj, ij$ が登場している。順序を保つように全ての添字を並べる方法は [p, q, i, j] または [p, i, q, j] があり、これに合うよう次元追加すれば転置不要で計算できる。（実装時は、最初から全ての添字を揃える必要は無く、二項演算の度に次元追加しても構わない）

r = Numo::DFloat[[0.8, 0.6], [-0.6, 0.8]]
a = Numo::DFloat[[1, 2], [3, 4]]

                                 #    p, q, i, j
r[true, :new, true, :new].shape  #=> [3, 1, 3, 1]
r[      true, :new, true].shape  #=>    [3, 1, 3]
a                        .shape  #=>       [3, 3]

# 数式通りの順序だと、途中で4次元配列ができてしまう
b = (r[true, :new, true, :new] * r[true, :new, true]).mulsum(a, -2, -1)

# （要素の演算が可換・結合的なら）縮約しやすい順序に変えて構わない
#  ij            q     i     j     j             p     q     i     i
b = a.mulsum(r[true, :new, true], -1).mulsum(r[true, :new, true], -1)

                                 #    p, i, q, j
r[true, true, :new, :new].shape  #=> [3, 3, 1, 1]
r                        .shape  #=>       [3, 3]
a[      true, :new, true].shape  #=>    [3, 1, 3]

b = a[true, :new, true].mulsum(r[true, true, :new, :new], -3).mulsum(r, -1)

一方で、この式は行列の積でも書き表せるが、掛ける順序に制約ができるうえに転置が必要になる： $b_{pq} = \sum_{i,j} r_{pi} a_{ij} r^T_{jq}$ 。

class Numo::NArray
	def matmul(other)
		self[false, :new].mulsum(other, -2)
	end
end

b = r.matmul(a).matmul(r.transpose(-1, -2))

3階以上でも問題ない。例えば $\sigma_{ij} = \sum_{k,l} d_{ijkl} \epsilon_{kl}$ （弾性係数テンソルをイメージした式）なら [i, j, k, l] で合わせればいい。1回の乗算で2つ以上の添字を縮約するときは、 #mulsum に指定する次元を増やす⁵。

d = Numo::Int64.new(3, 3, 3, 3).rand(10)
e = Numo::Int64.new(3, 3).rand(10)

         #    i, j, k, l
d.shape  #=> [3, 3, 3, 3]
e.shape  #=>       [3, 3]

s = d.mulsum(e, -2, -1)

レヴィ＝チヴィタ記号で遊んでみたりも。 $\sum_{j,k} \epsilon_{ijk} \epsilon_{ljk} = 2 \delta_{il}$ が成り立つか試してみる。

e = Numo::Int8.zeros(3, 3, 3)
e[0, 1, 2] = e[1, 2, 0] = e[2, 0, 1] = 1
e[0, 2, 1] = e[1, 0, 2] = e[2, 1, 0] = -1

# indices: i, l, j, k
e.expand_dims(-3).mulsum(e, -2, -1)
#=>
# Numo::Int8#shape=[3,3]
# [[2, 0, 0],
#  [0, 2, 0],
#  [0, 0, 2]]

定義済みの積

よく使う積に関しては、自分で実装しなくても Numo::NArray のメソッドとして存在している。内積を計算する箇所は実際に #mulsum を利用している。

ドキュメント： https://ruby-numo.github.io/narray/narray/Numo/NArray.html
ソースコード： https://github.com/ruby-numo/numo-narray/blob/v0.9.0.7/lib/numo/narray/extra.rb

#dot ：ドット積／行列積: 多次元配列の後ろ2次元を行列とみなし（1次元配列なら行or列ベクトル）、要素（行列）毎の積を計算する。なので普通の行列計算ならこのメソッドに任せるのが良さそう。
なお、線形代数ライブラリ Numo::Linalg が利用可能な場合はそちらへ投げている。
#inner ：内積: 多次元配列をベクトルの配列とみなし、ベクトル毎の内積を計算する。何次元目がベクトルかはキーワード引数で指定可能（デフォルトは -1 、つまり最後の次元）。
#outer ：外積（直積）: 多次元配列をベクトルの配列とみなし、ベクトル毎の外積を計算する。内積は次元が減る（ベクトル→スカラー）のに対し、こちらは次元が増える（ベクトル→2階テンソル）。
#kron ：クロネッカー積: テンソル積を計算し、各要素を多次元配列内に展開する。

それぞれ試してみた。

require 'numo/narray'

a = Numo::Int64[[1, 2], [3, 4]]
b = Numo::Int64[[1, 10], [100, 1000]]

a.dot(b)
#=>
# Numo::Int64#shape=[2,2]
# [[201, 2010],
#  [403, 4030]]

a.inner(b)
#=>
# Numo::Int64#shape=[2]
# [21, 4300]

# 参考：行列の積を内積で再現
a.expand_dims(-1).inner(b.expand_dims(-3), axis: -2)
#=>
# Numo::Int64#shape=[2,2]
# [[201, 2010],
#  [403, 4030]]

a.outer(b)
#=>
# Numo::Int64#shape=[2,2,2]
# [[[1, 10],
#   [2, 20]],
#  [[300, 3000],
#   [400, 4000]]]

a.kron(b)
#=>
# Numo::Int64#shape=[4,4]
# [[1, 10, 2, 20],
#  [100, 1000, 200, 2000],
#  [3, 30, 4, 40],
#  [300, 3000, 400, 4000]]

# 参考：テンソル積
a[false, *Array.new(b.ndim, :new)] * b
#=>
# Numo::Int64#shape=[2,2,2,2]
# [[[[1, 10],
#    [100, 1000]],
#   [[2, 20],
#    [200, 2000]]],
#  [[[3, 30],
#    [300, 3000]],
#   [[4, 40],
#    [400, 4000]]]]

次元毎でない指定方法も色々あるが、本記事では利用しないため省略する。 ↩
日本語の「外積」には複数の意味があるので注意。本記事ではクロス積の類は扱わない。 ↩
要素毎の計算は本来は配列形状が同じでないといけないが、大きさ1の次元にはbroadcastという規則が働く。 → ドキュメント ↩
例えばこちらを参照：行列の積 | 単位の密林 ↩
単純に可変長引数として与えることもできるし、キーワード引数なら axis: [-2, -1] のように配列で与えることもできる。 ↩

You get articles that match your needs
You can efficiently read back useful information
You can use dark theme

What you can do with signing up