TypeScriptでnumpyのような行列計算のライブラリを書いた話
簡単な自己紹介
odiak.icon 岩本海童 / odiak
ZOZO 研究所で研究者やってます
機械学習で計算した結果をブラウザでインタラクティブに可視化するために、TypeScriptでアプリケーションを作ってます
動機
具体的には、ほかの場所で計算した結果をJavaScriptを使ってブラウザで表示したい
表示のためにちょっとした計算が必要
データのサイズはそこまで大きくない
たかだか数万とかのオーダー
なので、計算は最適化されていなくてもよい
その代わり、numpyの便利な機能が使いたい
既存のライブラリ
einsumは本当に便利な機能で、特に多次元の配列を扱う際によく使う
じゃあちょっと作ってみるか？
方針
自分の欲しい機能だけ実装する
計算は速くなくてよい
dtypeみたいな配列の型の概念はなくていい
できたもの
いい名前が思いつかなかったのでnumjsにした
僕の知る限りnumjsというライブラリは3つある
どうせ自分しか使わないしいいや
実装した機能
インデックスによるアクセス
読み書き両方
書き込み時はarray.set([i, All], 0)という感じでnumpyでいうarray[i, :] = 0のようなワイルドカードによるアクセスもできる
reshape
transpose
swapAxes
次元の入れ替え
基本的な2項演算、1項演算(-)、基本的な数学関数
2項演算はブロードキャストされる
関数はまだexpだけ
スライス
配列の部分配列を返す
例: array.slice(0, All, NewAxis, range(1, 10))
Allはrange()のイディオム
集約
sum, meanなど
argMin, argMax
einsum
基本的にはnumpyのそれと同じ
einsum('i1,i2; i1,j,k -> i2,j,k', a, b)のように2文字以上のインデックス名を使えるようにした
numpyのそれは1文字の制約があるので、ときどきつらい
これを実装した後、Pythonでも同じようなインターフェースのeinsumを使うためのラッパーを作ってみた
clip
配列に含まれる値を指定した範囲に抑え込む
TypeScriptでライブラリを書くことについて
たぶん初めての体験
型があることの良さ
型があるとコードのミスに早いタイミングで気づけることが多い
型があるとAPIのインターフェースがきれいになる
例えばnumpyでは同じ書き方で書けるgetとsliceが別のAPIになっているのは、それぞれの戻り値の型が違うから
引数によって戻り値の方が変わるというのはバッドなパターンだと思う
書いているときも楽
複雑な型の値(例えばタプルの配列を受け取る関数の型)を扱うとき、自分でもどういう型の値だったのか忘れがちだが、型があると安心
失敗したらすぐ気づく
型を書くデメリットはほとんどないが、
若干面倒なときがある
推論が優秀なので、型を明示しなくてよいことが多いが、明示が必要なこともたまにある
TypeScriptのエラーがときどきわかりにくい
npmに公開するのは簡単
TSファイルを型定義ファイル付きでビルドするだけ
テスト
実装を間違えると厄介なのでテストを書いてみた
chaiのexpectを使ってmochaで実行する
小さいサイズの配列を用意して、計算結果を確認する
テストコードを見ると大体使い方がわかるかも
ドキュメントを書かないことへの言い訳
その他
VSCodeで書いた
tslintでコードを検査した
Prettierでコードをフォーマットした
感想
einsumとスライスとブロードキャストを実装するのが楽しかった
完全に自己満足の世界だけど
気になる人はコードを読んでみてほしい
TypeScriptの筋トレにもなった
機能追加の要望、プルリクエストなどは歓迎です