Your model is rather unusual in it’s definition. Without giving any feedback on whether it makes sense to write it down like that, you would need to make self.N and self.M, both pm.Data variables as well, which you update when doing predictions.