テンソルフローのブロードキャストドットプロダクト

Question

テンソルフローでは、次の問題があります。

I形状のテンソルM [BATCH_SIZE、dim_a、dim_b]及びマトリックスU形状の [BATCH_SIZE、dim_b]を有します。

M = tf.constant(shape=[batch_size, sequence_size, embed_dim]) 
U = tf.constant(shape=[batch_size, embed_dim]) 

は、私が何を達成するためには、私のバッチの各指標について[I、dim_a、dim_b]×[I、dim_b]の内積です。

P[i] = tf.matmul(M[i, :, :], tf.expand_dims(U[i, :], 1)) for each i. 

基本的には、バッチ軸上にドット積をブロードキャストする。これは可能ですか？どうすれば達成できますか？

Answer 1

これは）（tf.einsumで達成することができます：

import tensorflow as tf 
import numpy as np 

batch_size = 2 
sequence_size = 3 
embed_dim = 4 

M = tf.constant(range(batch_size * sequence_size * embed_dim), shape=[batch_size, sequence_size, embed_dim]) 
U = tf.constant(range(batch_size, embed_dim), shape=[batch_size, embed_dim]) 

prod = tf.einsum('bse,be->bs', M, U) 

with tf.Session(): 
    print "M" 
    print M.eval() 
    print 
    print "U" 
    print U.eval() 
    print 
    print "einsum result" 
    print prod.eval() 
    print 

    print "numpy, example 0" 
    print np.matmul(M.eval()[0], U.eval()[0]) 
    print 
    print "numpy, example 1" 
    print np.matmul(M.eval()[1], U.eval()[1]) 

出力：

M 
[[[ 0 1 2 3] 
    [ 4 5 6 7] 
    [ 8 9 10 11]] 

[[12 13 14 15] 
    [16 17 18 19] 
    [20 21 22 23]]] 

U 
[[2 3 3 3] 
[3 3 3 3]] 

einsum result 
[[ 18 62 106] 
[162 210 258]] 

numpy, example 0 
[ 18 62 106] 

numpy, example 1 
[162 210 258]