You can do this with the pymc_extras.statespace, using a combination of AR and Exogenous components. See an example here, though it’s a bit out of date and might need some updating.
I also want to add support for exogenous regressors directly in the BayesianSARMIA model, but haven’t gotten around to it yet. Help welcome
I’m currently working on state space models as part of a potential GSoC project, so this topic is quite relevant.
While my main focus would be on the Dynamic Factor Models (DFM), I might be able to look into AR processes with exogenous shocks as well, time permitting. If I have some bandwidth later in the summer, I’d be happy to explore this further or collaborate on it.