I have created a simplified example notebook where I show what I am trying to do:
I also started to study the blog post Out of model predictions with PyMC - PyMC Labs but I am still working on understanding what exactly is going on in the sections following “Making predictions on different models”.