[WIP] compute_expectation wrt plate and markov variables

[ordabayevy]

Addresses pyro-ppl/pyro#2724.

This function computes the expected value of cost terms wrt to log measures.

[ordabayevy]

@eb8680 can you review this?
Some notes:

• TraceMarkovEnum_ELBO that uses compute_expectation passes all of the tests locally for updated test_vectorized_markov::test_guide_enumerated_elbo from Make replay work correctly with vectorized_markov pyro#2726
• This PR can be simplified if funsor.adjoint can be used instead of forward_backward_terms function written here.

[eb8680]

@ordabayevy I'm impressed you got this working, and hopefully your model should be unblocked at least locally via this PR and pyro-ppl/pyro#2725 . Please let me know if that's not the case.

However, I am reluctant to merge this in its current form, especially without any tests - the algorithm implemented here is quite complex, and there is a substantial amount of code duplication from modified_partial_sum_product that I would prefer to reduce or avoid if possible given the changes we expect to make to the internals of modified_partial_sum_product in the future (#293). It's also not yet clear to me how we will need to change the algorithm or interface to facilitate sharing of computation across multiple integrands, as is done (somewhat imperfectly) in Pyro's pyro.infer.util.Dice.compute_expectation.

Can we talk sometime this week about how it might be simplified? I have some intuitions but not a fully fleshed out plan. It would probably be useful for both of us to prepare by reading this paper, which discusses general dynamic programming algorithms for computing expectations.

[ordabayevy]

Please let me know if that's not the case.

Okay, I'll let you know if there are any problems switching to Funsor backend and trying out this version of TraceMarkovEnum_ELBO.

Can we talk sometime this week about how it might be simplified?

Thursday and Friday will work for me.

[ordabayevy]

@eb8680 I've created an HMM version of my model where both m and theta are guide enumerated. I've tested it on simulated and real data and it seems to work! Although a bit slow which might be because compute_expectation works only in eager interpretation. As we discussed last week lazy implementation is blocked by the _scatter function which doesn't have a lazy equivalent right now.

[eb8680]

@ordabayevy that's great! I'll look into making _scatter a first-class funsor. One other thing you can try to speed it up is using the PyTorch JIT - it won't make the underlying tensor operations any faster but it should at least compile away the Funsor pattern-matching overhead.

That should be as simple as creating a JitTraceMarkovEnum_ELBO class similar to JitTraceEnum_ELBO and dropping it into your Pyro code:

from pyro.contrib.funsor.infer.elbo import Jit_ELBO
...
class JitTraceMarkovEnum_ELBO(Jit_ELBO, TraceMarkovEnum_ELBO):
    pass

For a Jit*_ELBO usage example see https://github.com/pyro-ppl/pyro/blob/dev/examples/contrib/funsor/hmm.py

[ordabayevy]

Closing as obsolete.