from theano.sandbox.rng_mrg import MRG_RandomStreams as RandomStreams
trng = RandomStreams(12345)
from theano import pp
from utils import *
rng = numpy.random.RandomState(1234)
x = T.vector()
pvals = SoftMax(x)
#y = trng.multinomial(pvals=pvals, dtype='int32')
#y = trng.multinomial_meanfield(pvals=pvals, dtype='int32')
y = trng.multinomial_straightthrough(pvals=pvals, dtype='int64')
emb = theano.shared(value=NormalInit(rng, 5, 1))
y_projection = T.dot(y, emb)
# Print un-optimized Theano graph
#print 'theano.printing.debugprint(y_projection)', theano.printing.debugprint(y_projection)
f = theano.function(
[x], [y_projection, pvals, y,
T.grad(T.sum(y_projection), pvals)])
#f = theano.function([x], [y_projection, pvals, y])
# Print optimized Theano graph
#print 'theano.printing.debugprint(f.maker.fgraph.outputs[0])', theano.printing.debugprint(f.maker.fgraph.outputs[0])
