@pascanur why do you use set_value? Why you don't put the subtensor in the graph of the inner fct? I think this is better and I think it would be faster. But I didn't time it. Make the inner take as input only pos and you can use the shared var for the input of the Op. Then you call set value only once on them outside the loop.

Also, after compiling the theano fct, so self.step.trust_input=True. This will lower the Theano overhead of processing the input. It this still take too much time, you can put pos as a shared variable, that get updated by the fct, so you only reset it before the loop.