numpy - 順次積と和を 3D でのより高速な行列演算に置き換える

Question

私の現在の theano スクリプトでは、ボトルネックは次のコードです。

import numpy as np

axis = 0
prob = np.random.random( ( 1, 1000, 50 ) )
cases = np.random.random( ( 1000, 1000, 50 ) )

start = time.time(  )
for i in xrange( 1000 ):
    result = ( cases * prob ).sum( axis=1-axis, keepdims=True )
print '3D naive method took {} seconds'.format( time.time() - start )
print result.shape
print

2D の場合、elementwise+sum を内積に置き換えると 5 倍のスピードアップが見られました。この場合、私を助けることができる行列演算はありますか?

編集：

Divakarは私にeinsumに基づくバージョンをくれました。ただし、私の意図はそれをtheanoに移植することであり、einsumはtheanoではサポートされていません。したがって、theanoに移植可能な代替手段は大歓迎です。