You are welcome! If the speed bottleneck persist and you are interested to explore solution in PyMC4/TFP, feel free to reach out and I can (or I will find someone who can ;-)) help you porting the model and inference. FWIW, the implementation of TFP SMC is very similar to PyMC3, with additional flexibility to use HMC as internal mutation which should make it scale much better to more dimensions. Also, if you are working on Covid and need to fit many time series at the same time, TFP generally gives better support for multi-batch which means you can fit multiple copies of the same model at once.