einsum
アインシュタインの縮約記法
テンソル積を書くと大量の$ \Sigmaが出てきてうざい
なので, $ \Sigmaを省略しよう, というのが事の始まり
$ \sum_{j} a_{i,j}b_{j,k} なら ij,jk → ik
プログラム的に考えれば, 単純にfor を外しただけ
例えば
ドット積
z[i,j] += x[i,j] * y[i,j]
np.einsum("ij,ij->ij", x, y)
内積
z[j,k] += x[j,i] * y[i,k]
np.einsum("ji,ik->jk", x, y)
軸の入れ替えなどでも使える
y = torch.einsum('nchw->nhwc', y).detach().cpu()
実際は遅いからやめたほうが良さそう
でもMAEの実装で使われてるからなぁ...