Theano version of a numpy einsum for two 3dim matrices

Question

I have two 3dim numpy matrices and I want to do a dot product according to one axis without using a loop in theano. a numpy solution with sample data would be like:

a=[ [[ 0, 0, 1, 1, 0,  0,  0,  0,  0,  0,  1,  0,  0,  1,  0],
  [ 1,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  1,  0,  1,  0],
  [ 0,  1,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  0,  1],
  [ 0,  1,  0,  0,  0,  0,  1,  0,  0,  0,  1,  0,  1,  0,  0]],
    [[ 0,  0,  1,  1,  0,  0,  0,  0,  0,  0,  1,  0,  0,  1,  0],
  [ 1,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  1,  0,  1,  0],
  [ 0,  1,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  0,  1],
  [ 0,  1,  0,  0,  0,  0,  1,  0,  0,  0,  1,  0,  1,  0,  0]],
 [ [ 0,  0,  1,  1,  0,  0,  0,  0,  0,  0,  1,  0,  0,  1,  0],
  [ 1,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  1,  0,  1,  0],
  [ 0,  1,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  0,  1],
  [ 0,  1,  0,  0,  0,  0,  1,  0,  0,  0,  1,  0,  1,  0,  0]],
 [ [ 0,  0,  1,  1,  0,  0,  0,  0,  0,  0,  1,  0,  0,  1,  0],
  [ 1,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  1,  0,  1,  0],
  [ 0,  1,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  0,  1],
  [ 0,  1,  0,  0,  0,  0,  1,  0,  0,  0,  1,  0,  1,  0,  0]],
 [[ 0,  0,  1,  1,  0,  0,  0,  0,  0,  0,  1,  0,  0,  1,  0],
  [ 1,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  1,  0,  1,  0],
  [ 0,  1,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  0,  1],
  [ 0,  1,  0,  0,  0,  0,  1,  0,  0,  0,  1,  0,  1,  0,  0]],
 [[ 0,  0,  1,  1,  0,  0,  0,  0,  0,  0,  1,  0,  0,  1,  0],
  [ 1,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  1,  0,  1,  0],
  [ 0,  1,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  0,  1],
  [ 0,  1,  0,  0,  0,  0,  1,  0,  0,  0,  1,  0,  1,  0,  0.]],
 [[ 0,  0,  1,  1,  0,  0,  0,  0,  0,  0,  1,  0,  0,  1,  0],
  [ 1,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  1,  0,  1,  0],
  [ 0,  1,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  0,  1],
  [ 0,  1,  0,  0,  0,  0,  1,  0,  0,  0,  1,  0,  1,  0,  0]]]

b=[[[ 0,  0,  1,  0,  0.],
  [ 1,  0,  0,  0,  0.],
  [ 0,  0,  0,  0,  0.],
  [ 0,  1,  0,  0,  0.]],
 [[ 0,  0,  1,  0,  0.],
  [ 1,  0,  0,  0,  0.],
  [ 0,  0,  0,  0,  0.],
  [ 0,  1,  0,  0,  0.]],
 [[ 0,  0,  1,  0,  0.],
  [ 1,  0,  0,  0,  0.],
  [ 0,  0,  0,  0,  0.],
  [ 0,  1,  0,  0,  0.]],
 [[ 0,  0,  1,  0,  0.],
  [ 1,  0,  0,  0,  0.],
  [ 0,  0,  0,  0,  0.],
  [ 0,  1,  0,  0,  0.]],
 [[ 0,  0,  1,  0,  0.],
  [ 1,  0,  0,  0,  0.],
  [ 0,  0,  0,  0,  0.],
  [ 0,  1,  0,  0,  0.]],
 [[ 0,  0,  1,  0,  0.],
  [ 1,  0,  0,  0,  0.],
  [ 0,  0,  0,  0,  0.],
  [ 0,  1,  0,  0,  0.]],
 [[ 0,  0,  1,  0,  0.],
  [ 1,  0,  0,  0,  0.],
  [ 0,  0,  0,  0,  0.],
  [ 0,  1,  0,  0,  0.]]]
dt = np.dtype(np.float32)
a=np.asarray(a,dtype=dt)
b=np.asarray(b,dtype=dt)
print(a.shape)
print(b.shape)

where "a", has the shape of (7, 4, 15) and "b", has the shape of (7, 4, 5). "c", is defined as dot product of "a" and "b":

c = np.einsum('ijk,ijl->ilk',a,b)

I am looking for a theano implementation of this example to calculate "c".

Any ideas?

Any particular reason you want `theano` over einsum (speed/memory usage/etc)? — Daniel, Nov 30 '15 at 19:55
@Ophion of course :) I need theano for gradient calculation. This will (hopefully) be part of a deep neural net. — PickleRick, Nov 30 '15 at 20:13
Well, in principle [batched_tensordot](http://deeplearning.net/software/theano/library/tensor/basic.html#theano.tensor.batched_tensordot) is exactly what you need. Although, I cannot get it to work. — Daniel, Nov 30 '15 at 20:28
@Ophion I thoughts so, but I also couldn't figure it out. tc = T.batched_tensordot(ta,tb, axes=0) f_c = theano.function(inputs=[ta,tb], outputs=tc) print(np.shape( f_c(a,b))) gives me an output of dimension (7, 4, 5, 4, 15) which is wrong. The other axes don't work. using theano.tensor.batched_dot: f_c = theano.function(inputs=[ta,tb], outputs=tc) print(np.shape( f_c(a.transpose([1,0,2]),b.transpose([0,1,2])))) I tried a couple of transpose shapes, but it didn't work and raised an exception of dimension mismatch. — PickleRick, Nov 30 '15 at 20:49
I got a compile error with mine. The call should be: `T.batch_tensor, tb, ta, axis=[[1],[1]])`. Have not used theano much though. — Daniel, Nov 30 '15 at 23:02
@Ophion great! this one works :) `tc = T.batched_tensordot(ta, tb, axes=[[1],[1]]) f_c= theano.function(inputs=[ta,tb], outputs=tc) print(np.shape( f_c(a,b)))` thank's! — PickleRick, Nov 30 '15 at 23:17

score 3 · Accepted Answer · edited Nov 30 '15 at 23:55

3

To finish out this question:

import theano as th
import then.Tensor as T

ta = T.tensor3('a')
tb = T.tensor3('b')

tc = T.batched_tensordot(ta, tb, axes=[[1],[1]])

......

edited Nov 30 '15 at 23:55

PickleRick

419
1
5
13

answered Nov 30 '15 at 23:21

Daniel

19,179
7
60
74

do you know a similar solution to divide instead of dot product? something that works in a similar way to divide instead of multiply? – PickleRick Dec 03 '15 at 13:07
Like `c[i,j] = np.sum(a[i,:]/b[:,j])`? No idea in this case. If you want general broadcasting techniques I would need to see an example. – Daniel Dec 04 '15 at 13:56
this does the job: a[:, :, None] / b – PickleRick Dec 04 '15 at 13:58

Theano version of a numpy einsum for two 3dim matrices

1 Answers1

Linked