Custom discrete inference in SVI

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.

have you tried using TraceMarkovEnum_ELBO like here

What is ‘here’ referring to? Can you repost the link?