Swift for TensorFlow の randomUniform を読む
モチベーション
以下の Swift for TensorFlow のスレッドで Chris が発言した
https://groups.google.com/a/tensorflow.org/forum/#!topic/swift/IQcQIOslt-I
Take a look at Tensor.swift in the TensorFlow module. We implemented this with the standard swift generics system, no special support necessary.
をみて少し気になった
また前回書いたswift-models の MNIST.swift を読むでも登場したrandomUniformに関して、どういった実装をしているのか気になり読んでみた
https://github.com/tensorflow/swift-models/blob/master/MNIST/MNIST.swift#L61-L62
読んでいくうちに結構面白いと自分でも思えたので過程をブログにすることにした
randomUniform とは
エントリポイントとなるのは以下の initializer
https://github.com/google/swift/blob/tensorflow/stdlib/public/TensorFlow/Tensor.swift#L533-L546
与えられた TensorShape の要素を全て0-1の正規分布の範囲内の値で埋め尽くすというもの
まずここで気づくのが、単純に Swift for TensorFlow が TensorFlow と Swift のブリッジング/マッパーではないということで、結構な部分（特に型まわり）が pure Swift で実装されている
また読み進める中でいくつかの _ 付きの @ attributes がある
@_fixed_layout @_versioned @_inlineable @inline(__always)
しっかり調べられていないので後で調べてまた別でブログにしたい
https://bugs.swift.org/browse/SR-260 https://stackoverflow.com/questions/46524232/access-control-in-swift-4 https://lists.swift.org/pipermail/swift-evolution/Week-of-Mon-20171002/040061.html
ここからスタート、基本的に見たことのないユニットは一つずつあげていく
TensorShape
https://github.com/google/swift/blob/tensorflow/stdlib/public/TensorFlow/TensorShape.swift
配列の薄いラッパーで、次元を表現するときに使う
ExpressibleByArrayLiteral に対応しているので利用する側は配列と思って使っていても問題ない
RandomState
https://github.com/google/swift/blob/tensorflow/stdlib/public/TensorFlow/Utilities.swift#L148
擬似乱数を表現するクラス、全体的にデフォルトだと RandomState.global を使用しているところが多い
Tensor.init(randomUniform shape: TensorShape, state: RandomState? = nil)
ここでようやく元の initializer に戻れる
この initializer の実態は以下の initializer のラッパーであることがわかる
https://github.com/google/swift/blob/tensorflow/stdlib/public/TensorFlow/Tensor.swift#L543-L545
ラップされた以下の initializer を RAND_MAX で割っていることから、以下の initializer は整数値の range での正規分布を元にした random ということが推測できる
ここで登場している Scalar クラスはどこで定義されているのか分からなかった。が、行列計算におけるスカラー値を表現していることは容易にわかる
また以下にあるように、AccelerableByTensorFlow プロトコルに準拠している型は Scalar として使えるようだった
https://github.com/google/swift/blob/tensorflow/stdlib/public/TensorFlow/DataTypes.swift#L34
init(randomStandardUniform shape: TensorShape, state: RandomState? = nil)
次にラップされていた上記を読み進める前に、いくつか調べておかないといけないユニットがある
_TFHoistable
_TFHoistable は () -> TensorHandle<Scalar> な closure を受け取り、その closure の返り値を return するだけのもの
https://github.com/google/swift/blob/tensorflow/stdlib/public/TensorFlow/Tensor.swift#L124
TensorHandle
TensorHandle は Tensor の核となるもので、実は Tensor は TensorHandle のラッパー
TensorHandle は ops や #tfop() を実行するための型
実際に TensorFlow とやりとりするときに、パラメータの型をより静的にコンパイラに伝えるためのものだと思った
_TFTensorFromScalars<Scalar>(_ scalars: [Scalar], shape: [Int32])
https://github.com/google/swift/blob/tensorflow/stdlib/public/TensorFlow/Tensor.swift#L94-L107
内部の実装を読むと渡された scalar の配列と shape のサイズに違いがないかチェックし、
TensorHandle を初期化している。 scalarsInitializer では TensorHandle に渡す scalars を全てアロケートしないといけないらしい。これも TensorFlow のコアに渡すためなんだろうか
https://github.com/google/swift/blob/tensorflow/stdlib/public/TensorFlow/TensorHandle.swift#L45
TensorFlow.init(randomStandardUniform shape: TensorShape, state: RandomState? = nil)
https://github.com/google/swift/blob/tensorflow/stdlib/public/TensorFlow/Tensor.swift#L521-L529
ここまででようやくメインに戻ってくることができる
やっていることとしては、Tensor を TensorHandle を引数にして初期化している
その TensorHandle は _TFHoistable 経由で作成しており、その closure の内部で _TFTensorFromScalars を使って、ランダムに作られた Int の配列から TensorHandle を作っていることがわかる
まとめ
サンプルで使われていたある関数を起点として、いかに静的型付けと TensorFlow のコアとのブリッジングを成立させるかの足がかりが見えた。
今回見えた TensorFlow/DataTypes.swift や TensorFlow/TensorHandle.swift などはそれに強く関わる部分だと思うので、もっと調べる
https://github.com/google/swift/blob/tensorflow/stdlib/public/TensorFlow/DataTypes.swift https://github.com/google/swift/blob/tensorflow/stdlib/public/TensorFlow/TensorHandle.swift
