Have you seen/tried the statespace VARMAX implementation? It has helpers for forecasts and IRFs that automate these issues.
Regardless of how you fit the model I think the easiest way to handle these post-estimation tasks is to re-cast the model in statespace form and just iterate on x_t = T x_{t-1} + d + \varepsilon_t. That’s what I do in this blog post, which doesn’t use PyMC-Statespace. It’s much easier to use an auxiliary forecasting model than to try to wrestle with the xarrays. Here’s an example where I wrestled with the vectorized ufuncs but then switched to a forecasting model for an easier/happier life.