No you can’t write the PyMC model with a loop. You can perhaps write the scan function in JAX and wrap it in a Pytensor Op as a quick alternative. There are some recipes on how to do something like that here: How to use JAX ODEs and Neural Networks in PyMC - PyMC Labs