I would like to combine pyro’s SVI with a custom discrete inference step for a factorial HMM model with a rather complex likelihood. We have tried to use pyro’s built in discrete inference but find it too slow. My feeling is that I could condition my model function on either fixed states (viterbi), or a categorical over states (ie marginals, forward backward), take a step for the non-discrete model parameters. Then I would need to implement an update step for the discrete variables given the current guide and likelihood. Is there a tractable way of implementing such a step within pyro? How does one obtain the expected likelihood of the non-discrete variables with respect to the variational approximation (guide), as I think I would need these for the discrete state updates.