Search code examples
pythonnumpydeep-learningtheanonumpy-einsum

how to use batch_tensordot in theano like numpy.einsum

I have a tensor3 with shape (3, 4, 5) and another tensor4 with shape (3, 4, 7, 5). In numpy,

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

but in theano , how to do it .

Solution

  • The first step is to transpose and reshape your tensor so that only the first dimension gets preserved. In that case it is quite simple, you just have to combine the first two dimensions:

    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))

    Then, you specify to batched_tensordot that you are going to sum axis 1 of x_ with axis 2 of y_:

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

    Finally, reshape z_ to get the first two dimensions:

    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)