Search code examples
numpytheanodeep-learning

theano - use tensordot compute dot product of two tensor


I want to use tensordot to compute the dot product of a specific dim of two tensors. Like:

A is a tensor, whose shape is (3, 4, 5) B is a tensor, whose shape is (3, 5)

I want to do a dot use A's third dim and B's second dim, and get a output whose dims is (3, 4)

Like below:

for i in range(3):
    C[i] = dot(A[i], B[i])

How to do it by tensordot?


Solution

  • Well, do you want this in numpy or in Theano? In the case, where, as you state, you would like to contract axis 3 of A against axis 2 of B, both are straightforward:

    import numpy as np
    
    a = np.arange(3 * 4 * 5).reshape(3, 4, 5).astype('float32')
    b = np.arange(3 * 5).reshape(3, 5).astype('float32')
    
    result = a.dot(b.T)
    

    in Theano this writes as

    import theano.tensor as T
    
    A = T.ftensor3()
    B = T.fmatrix()
    
    out = A.dot(B.T)
    
    out.eval({A: a, B: b})
    

    however, the output then is of shape (3, 4, 3). Since you seem to want an output of shape (3, 4), the numpy alternative uses einsum, like so

    einsum_out = np.einsum('ijk, ik -> ij', a, b)
    

    However, einsum does not exist in Theano. So the specific case here can be emulated as follows

    out = (a * b[:, np.newaxis]).sum(2)
    

    which can also be written in Theano

    out = (A * B.dimshuffle(0, 'x', 1)).sum(2)
    out.eval({A: a, B: b})