Description
The most requested feature for the statespace module is to handle multiple time series in a single MCMC run. This will require support for batch dimensions. I think the easiest way to attack this will be to refactor the KalmanFilter class to be an OpFromGraph. My original class-based design was inspired by the statsmodels implementation, but it doesn't take full advantage of pytensor.
Thinking more about this, I don't know if KF needs to be an OFG as a first step. It might still be nice to have a AbstractKalmanFilter dummy that we can rewrite to specialized cases, but when I wrote this issue I was a bit obsessed with OFG. The custom gradients are still on my mind, though. So the next two sentences remain true:
An additional advantage of this will be the ability to define a custom gradient. See #332.
Finally, it will let us handle special case filters via rewrites, rather than asking the user to pick a filter up front.