import numpy as np
import theano
import theano.tensor as T
from theano.tensor.shared_randomstreams import RandomStreams

# Example of using scan within scan in Theano

numpy_rng = np.random.RandomState(1234).randint(2 ** 30)

D = T.fvector("D")
X = T.fmatrix("X")

def test(X1,D1):
    X_hat=T.nnet.softmax(X1)
    components, updates = theano.scan(fn=lambda x1,d: RandomStreams(numpy_rng).multinomial(size=(1,1), n=d, pvals=x1, dtype=theano.config.floatX),
                                                            outputs_info=None, sequences=[X_hat,D1])
    return [[T.squeeze(components), D1+1],updates] # pass updates as second element in a list. If you have more updates (e.g. more loops) here,
                                                   # concatenate them instead of creating the list (i.e. pass updates1+updates2 instead of [update1,update2] as in case of output vars

c, up = theano.scan(fn=lambda dd,xx: test(xx,dd), outputs_info=[None, D],non_sequences=X,n_steps=3)

f=theano.function([X,D],c,updates=up,allow_input_downcast=True)

c=f(np.random.rand(3,5),np.array([15,5,10]))
print c[0]