Repository: pyro-ppl/pyro

Issue:
[bug] Sign error in JitTraceGraph_ELBO
Trying out this class and ended up with a rather strange looking loss curve (I was using a LR scheduler which explains the waviness):

![elbo_bug](https://user-images.githubusercontent.com/963647/113894124-b22ac200-9795-11eb-84ee-c2c0fab6453c.png)

It looks like the training losses have a sign error, although they are still optimizing in the right direction. This bug doesn't appear in `TraceGraph_ELBO` or `JitTrace_ELBO`.

Looking at the code, I think this is just because the JITed class overrides `loss_and_grads` and lacks a negative sign which appears in the non-JIT version. The validation curve is generated via `SVI.evalulate_loss` which calls `TraceGraph_ELBO.loss`  and includes the negation.

So it seems like a simple fix, but before I made a PR I wanted to flag the issue in case I'm missing something. There seem to be negations in a couple places and it seems easy to mess up

Added tests for JIT Elbo classes in test_valid_models
I was attempting to fix #2798 by simply adding a minus sign to the `JitTraceGraph_ELBO` class and discovered some issues while testing. As I cannot figure out what is going on, @fritzo suggested I open a PR with failing tests for someone else to look at.

So, for every use of an ELBO in `test_valid_models`, I added another test that used the Jit version of that ELBO if there was one. For a few tests this meant introducing parametrization (the TraceMeanField_ELBO tests lived on their own, and now they have a parameter for the normal and Jit class). Also, in cases where the test used `if Elbo is TraceEnum_ELBO` I changed it to `if issubclass(Elbo, TraceEnum_ELBO)` so that the Jit version is included.

This adds more than 200 failing tests (and some passing ones as well). 😬  I suspect that many of them are simply because the test model is not compatible with JITting, but it also uncovers a number where the outputs simply don't match the expected value from `JitTraceGraph_ELBO` and `JitTraceMeanField_ELBO`.


Discussion:
Good catch @jamestwebber!  We had a similar bug #1982 and added what I thought was a thorough regression test in #1988.  I think the first step in a PR to fix your bug would be a similar regression test (say in `test_valid_models.py`) that fails for the current `JitTraceGraph_ELBO` implementation.
Ah yes, digging into the existing tests (just running `test_plate_ok`) it looks like the loss ends up at 0.0 and so -0.0 passes just fine. (oh, and it looks like the previous fix just missed the JIT version)

I'll write a PR in a bit.
Facing the same issue with Trace_ELBO(), Beta distribution and MNIST data. Any advances on this?
Hmm I guess @jamestwebber added failing tests in #2799, but we haven't gotten around to making all those pass. I can add this to my TODO list


Resolving patch:
diff --git a/pyro/infer/trace_mean_field_elbo.py b/pyro/infer/trace_mean_field_elbo.py
--- a/pyro/infer/trace_mean_field_elbo.py
+++ b/pyro/infer/trace_mean_field_elbo.py
@@ -120,7 +120,10 @@ def _differentiable_loss_particle(self, model_trace, guide_trace):
                             kl_qp, scale=guide_site["scale"], mask=guide_site["mask"]
                         )
                         if torch.is_tensor(kl_qp):
-                            assert kl_qp.shape == guide_site["fn"].batch_shape
+                            assert (
+                                torch._C._get_tracing_state()
+                                or kl_qp.shape == guide_site["fn"].batch_shape
+                            )
                             kl_qp_sum = kl_qp.sum()
                         else:
                             kl_qp_sum = (
diff --git a/pyro/infer/tracegraph_elbo.py b/pyro/infer/tracegraph_elbo.py
--- a/pyro/infer/tracegraph_elbo.py
+++ b/pyro/infer/tracegraph_elbo.py
@@ -412,12 +412,12 @@ def jit_loss_and_surrogate_loss(*args, **kwargs):
 
             self._jit_loss_and_surrogate_loss = jit_loss_and_surrogate_loss
 
-        loss, surrogate_loss = self._jit_loss_and_surrogate_loss(*args, **kwargs)
+        elbo, surrogate_loss = self._jit_loss_and_surrogate_loss(*args, **kwargs)
 
         surrogate_loss.backward(
             retain_graph=self.retain_graph
         )  # triggers jit compilation
 
-        loss = loss.item()
+        loss = -elbo.item()
         warn_if_nan(loss, "loss")
         return loss
diff --git a/pyro/ops/provenance.py b/pyro/ops/provenance.py
--- a/pyro/ops/provenance.py
+++ b/pyro/ops/provenance.py
@@ -43,6 +43,9 @@ class ProvenanceTensor(torch.Tensor):
     :param frozenset provenance: An initial provenance set.
     """
 
+    _t: torch.Tensor
+    _provenance: frozenset
+
     def __new__(cls, data: torch.Tensor, provenance=frozenset(), **kwargs):
         assert not isinstance(data, ProvenanceTensor)
         if not provenance:


Test patch:
diff --git a/tests/infer/test_jit.py b/tests/infer/test_jit.py
--- a/tests/infer/test_jit.py
+++ b/tests/infer/test_jit.py
@@ -24,12 +24,13 @@
     TraceEnum_ELBO,
     TraceGraph_ELBO,
     TraceMeanField_ELBO,
+    TraceTMC_ELBO,
     infer_discrete,
 )
 from pyro.optim import Adam
 from pyro.poutine.indep_messenger import CondIndepStackFrame
 from pyro.util import ignore_jit_warnings
-from tests.common import assert_equal
+from tests.common import assert_close, assert_equal
 
 
 def constant(*args, **kwargs):
@@ -251,6 +252,53 @@ def f(probs):
     assert log_prob.shape == shape[-1:] + batch_shape
 
 
+@pytest.mark.parametrize(
+    "Elbo",
+    [
+        Trace_ELBO,
+        JitTrace_ELBO,
+        TraceGraph_ELBO,
+        JitTraceGraph_ELBO,
+        TraceEnum_ELBO,
+        JitTraceEnum_ELBO,
+        TraceMeanField_ELBO,
+        JitTraceMeanField_ELBO,
+        TraceTMC_ELBO,
+    ],
+)
+def test_loss(Elbo):
+    pyro.clear_param_store()
+    data = torch.tensor(1.0)
+
+    def model(data):
+        loc = pyro.sample("loc", dist.Normal(0, 1))
+        scale = pyro.sample("scale", dist.LogNormal(0, 1))
+        pyro.sample("obs", dist.Normal(loc, scale), obs=data)
+
+    def guide(data):
+        loc_loc = pyro.param("loc_loc", lambda: torch.tensor(1.0))
+        scale_loc = pyro.param("scale_loc", lambda: torch.tensor(-1.0))
+        pyro.sample("loc", dist.Normal(loc_loc, 2))
+        pyro.sample("scale", dist.LogNormal(scale_loc, 0.1))
+
+    elbo = Elbo(
+        num_particles=10_000,
+        vectorize_particles=True,
+        max_plate_nesting=0,
+        strict_enumeration_warning=False,
+    )
+    expected = 18.611
+
+    try:
+        actual = elbo.loss(model, guide, data)
+        assert_close(actual, expected, rtol=0.1)
+    except NotImplementedError:
+        pass
+
+    actual = elbo.loss_and_grads(model, guide, data)
+    assert_close(actual, expected, rtol=0.1)
+
+
 @pytest.mark.parametrize("num_particles", [1, 10])
 @pytest.mark.parametrize(
     "Elbo",