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])