batched_dot.py

import timeit

import numpy as np
import theano
import theano.tensor as T

x = theano.shared(np.random.rand(500, 100).astype(theano.config.floatX))
W = theano.shared(np.random.rand(10, 100, 50).astype(theano.config.floatX))
y = theano.shared(np.random.randint(0, 10, 500).astype(theano.config.floatX))
y = T.cast(y, 'int32')

def batched_dot1():
    return T.batched_dot(x, W[y])

def batched_dot2():
    return (x[...,np.newaxis] * W[y]).sum(axis=1)

def batched_dot3():
    output, _ = theano.scan(
        fn=lambda x, y, W: T.dot(x, W[y]),
        sequences=[x, y],
        non_sequences=W)
    return output


for fn in (batched_dot1, batched_dot2, batched_dot3):
    print("Running", fn.__name__)
    f = theano.function([], fn())
    print("Took {} seconds".format(
        timeit.timeit('f()', setup='from __main__ import f', number=5000)))