theano_assertion_error.py

from theano import tensor as T 
import theano
import numpy as np

floatX = 'float32'
num_blocks = 2
block_dims = 3
hidden_dims = 5
output_dims = 7

if __name__ == '__main__':
    randm = lambda *dims: np.random.uniform(low=-0.01, high=0.01, size=dims)
    E = [theano.shared(np.asarray(randm(100, block_dims), dtype=floatX))]
    W0_values = np.asarray(randm(num_blocks, hidden_dims, block_dims), dtype=floatX)
    W0 = theano.shared(W0_values, 'W0')
    W1_values = np.asarray(randm(hidden_dims, output_dims), dtype=floatX)
    W1 = theano.shared(W1_values, 'W1')
    
    x = T.matrix('x', 'int64')
    y_hat = T.matrix('y_hat', 'int64')
    proj = T.stack(*[E[i][x[:,i]] for i in range(len(E))])
    proj.name = 'proj'
    hidden = T.prod(T.batched_tensordot(proj, W0, [2, 2]), axis=0)
    hidden.name = 'hidden'
    y = T.nnet.softmax(T.dot(hidden, W1))
    y.name = 'y'
    params = E + [W0, W1]
    cost = T.mean(-T.log(y[T.arange(y_hat.shape[0]), y_hat[:,0]]))
    grad = T.grad(cost, params)
    theano.function([x, y_hat], grad)