python - numpy.einsum のような theano で batch_tensordot を使用する方法

Question

形状 (3, 4, 5) の tensor3 と形状 (3, 4, 7, 5) の別の tensor4 があります。ナンパでは、

 result = np.einsum("ijk, ijmk->ijm", tensor3, tensor4)
 print result.shape 
 (3, 4, 7)

しかし、テアノでは、それを行う方法。

Answer 1

最初のステップは、最初の次元のみが保持されるように、テンソルを転置して再形成することです。その場合、最初の 2 つの次元を結合するだけです。

x = tensor.tensor3()
y = tensor.tensor4()
i, j, m, k = y.shape

x_ = x.reshape((i * j, k))
y_ = y.reshape((i * j, m, k))

次に、の軸 1と の軸 2batched_tensordotを合計するように指定します。x_y_

z_ = tensor.batched_tensordot(x_, y_, (1, 2))  # shape (i * j, m)

最後に、形状を変更z_して最初の 2 つの次元を取得します。

z = z_.reshape((i, j, m))
print(z.eval({x: np.zeros((3, 4, 5)), y: np.zeros((3, 4, 7, 5))}).shape)
# (3, 4, 7)