Add analytic single-sample test for ELBO gradient estimators

[ordabayevy]

This test compares analytically derived pathwise and score function gradient estimators for a single particle to gradients obtained by Trace_ELBO and TraceEnum_ELBO. When reparameterization is disabled gradients obtained by Trace_ELBO don't match analytically derived score function gradients.

By comparing gradient values it seems like grads obtained by Trace_ELBO are missing a second term (delbo_dscale and delbo_dloc) in the derivative of a surrogate loss:

expected_grads = {
    "scale": -(dlogq_dscale * elbo + delbo_dscale).sum(0, keepdim=True),
    "loc": -(dlogq_dloc * elbo + delbo_dloc),
}
# actual_grads are missing delbo_dscale term in "scale" and delbo_dloc term in "loc"

On the other hand, TraceEnum_ELBO passes all tests.

[fritzo]

From this test alone it is unclear whether there is a bug in stochastic gradient estimation; rather it is clear only that Trace_ELBO disagrees with a particular hand-coded estimator (at a single sample). That is, it still seems possible that Trace_ELBO might agree with your hand-coded estimator in expectation, but disagree at each sample. To confirm this is a true inference bug, we could also try setting num_particles=100000, vectorize_particles=True and see whether the expected gradients coming out of Trace_ELBO and TraceEnum_ELBO agree. WDYT?

[ordabayevy]

You are right! The missing term is a score function and its expectation is zero (which is cleverly removed by Trace_ELBO). So the expected gradients of Trace_ELBO and TraceEnum_ELBO agree (as tested in test_subsample_gradient).

I fixed the hand-coded formula for Trace_ELBO so the test should pass now. @fritzo do you want to add the test to the repo? If not, I'll close the PR.