Hamiltonian Monte Carlo with Energy Conserving Subsampling

[OlaRonning]

Introduces HMCECS to numpyro.

This PR includes:

• Taylor proxy [section 4.2];
• Difference estimation method [section 4.1].

[fehiepsi]

Can you come with some suggestions for test cases?

I think we can test if taylor_proxy returns correct results and estimate_likelihood returns values that are close to the true loglikelihood with taylor_proxy (in the sense that variance is small if we change subsample indices and latent values). I think Normal-Normal model is a good test case but we can increase the latent dimensions to a few. What do you think?

[OlaRonning]

@fehiepsi those tests sound good. I'll add them in.

[OlaRonning]

@fehiepsi should the python version in CI be bumped to 3.9? looks like it's supported by Github Actions.

[fehiepsi]

should the python version in CI be bumped to 3.9

It is not needed. We typically want to test for the oldest Python version that we support.

[fehiepsi]

Looks clean to me! Thanks for addressing my comments, @OlaRonning! Please let me know if you want to add more tests in another PR or if this is ready to merge when the following comments are addressed.

[OlaRonning]

Hi @fehiepsi, thanks for the reviews!

I'll add the test cases (proxy result and low variance) you suggested and then it's ready to be merged.