I think you are on the right path - you just need to make sure the tt.scan part is implemented correctly (which is the most challenging part!). Have a look at eg: https://docs.pymc.io/notebooks/PyMC3_tips_and_heuristic.html#PyMC3-implementation-using-theano.scan about how to implement a tt.scan and compare with python for loop