[theano-users] broadcast_batched_dot

[luke]

Hi all,


I want to achieve a "broadcast batched dot" operation in theano, such that
the two arguments A and B with shapes

A.shape = [2,5,7,3]
B.shape = [5,3,6]


produce an output C of shape tensor4 [2,5,7,6], with a np equivalent of:

for i in range(A.shape[0]):
for j in range(A.shape[1]):
C[i,j,:,:] = np.dot( A[i,j,:,:], B[j,:,:] )


So, basically, the last two dimensions of A and B are multiplied together
with dot, dimension 1 of A and 0 of B are batched, and dimension 0 of A is
broadcasted onto B.
I've played around a bit with T.batched_tensordot, but could not achieve
this.

The only way I could make this work involves a scan over dimension 0 of A,
and a T.batched_dot over the remaining 3 dimensions. But this is of course
dauntingly slow.


Any ideas?


br,
Luke

[Pascal Lamblin]

Hi,

You can use batched_tensordot for that, but it assumes that the
"batched" dimension is the first one, so you'd have to transpose A first
so that the "5" is first, and then transpose the result back to get C.

So here, you'd do something like:
A_T = A.transpose(1, 0, 2, 3) # shape = [5, 2, 7, 3]
C_T = T.batched_tensordot(A_T, B, axes=[3, 1]) # "axes" matches the "3"
between A_T and B, shape = [5, 2, 7, 6]
C = C_T.transpose(1, 0, 2, 3) # shape = [2, 5, 7, 6]

Post by luke

C.eval({A: np.zeros((2, 5, 7, 3)), B: np.ones((5, 3, 6))}).shape

(2, 5, 7, 6)

Post by luke
Hi all,
I want to achieve a "broadcast batched dot" operation in theano, such
that the two arguments A and B with shapes
A.shape = [2,5,7,3]
B.shape = [5,3,6]
            C[i,j,:,:] = np.dot( A[i,j,:,:], B[j,:,:] )
So, basically, the last two dimensions of A and B are multiplied
together with dot, dimension 1 of A and 0 of B are batched, and
dimension 0 of A is broadcasted onto B.
I've played around a bit with T.batched_tensordot, but could not achieve
this.
The only way I could make this work involves a scan over dimension 0 of
A, and a T.batched_dot over the remaining 3 dimensions. But this is of
course dauntingly slow.
Any ideas?
br,
Luke
--
---
You received this message because you are subscribed to the Google
Groups "theano-users" group.
To unsubscribe from this group and stop receiving emails from it, send
For more options, visit https://groups.google.com/d/optout.