Repository: pyro-ppl/pyro

Issue:
[FR] Default to max_plate_nesting="auto"
@mbabadi requested an easier-to-use interface for `TraceEnum_ELBO`, where `max_plate_nesting` could be inferred automatically.

One possible solution would be to set the default to "auto" rather than `float('inf')`, and use `EscapeMessenger` to optimistically start at zero and then on encountering a `plate` exceeding headroom: break, increment, and retry. This should work for most models and will have only `O(1)` overhead overall since it should only happen at the first SVI step. Ideally `max_plate_nesting` would almost never need to be specified.


Discussion:

Resolving patch:
diff --git a/pyro/infer/elbo.py b/pyro/infer/elbo.py
--- a/pyro/infer/elbo.py
+++ b/pyro/infer/elbo.py
@@ -1,11 +1,16 @@
 from __future__ import absolute_import, division, print_function
 
+import logging
 import warnings
 from abc import ABCMeta, abstractmethod
 
 from six import add_metaclass
 
 import pyro
+import pyro.poutine as poutine
+from pyro.infer.util import is_validation_enabled
+from pyro.poutine.util import prune_subsample_sites
+from pyro.util import check_site_shape
 
 
 @add_metaclass(ABCMeta)
@@ -23,19 +28,20 @@ class ELBO(object):
     :param num_particles: The number of particles/samples used to form the ELBO
         (gradient) estimators.
     :param int max_plate_nesting: Optional bound on max number of nested
-        :func:`pyro.plate` contexts. This is only required to enumerate over
-        sample sites in parallel, e.g. if a site sets
-        ``infer={"enumerate": "parallel"}``.
+        :func:`pyro.plate` contexts. This is only required when enumerating
+        over sample sites in parallel, e.g. if a site sets
+        ``infer={"enumerate": "parallel"}``. If omitted, ELBO may guess a valid
+        value by running the (model,guide) pair once, however this guess may
+        be incorrect if model or guide structure is dynamic.
     :param bool vectorize_particles: Whether to vectorize the ELBO computation
         over `num_particles`. Defaults to False. This requires static structure
-        in model and guide. In addition, this requires specifying a finite
-        value for `max_plate_nesting`.
+        in model and guide.
     :param bool strict_enumeration_warning: Whether to warn about possible
         misuse of enumeration, i.e. that
         :class:`pyro.infer.traceenum_elbo.TraceEnum_ELBO` is used iff there
         are enumerated sample sites.
-    :param bool retain_graph: Whether to retain autograd graph during an SVI step.
-        Defaults to None (False).
+    :param bool retain_graph: Whether to retain autograd graph during an SVI
+        step. Defaults to None (False).
 
     References
 
@@ -62,14 +68,47 @@ def __init__(self,
         self.max_plate_nesting = max_plate_nesting
         self.vectorize_particles = vectorize_particles
         self.retain_graph = retain_graph
-        if self.vectorize_particles:
-            if self.num_particles > 1:
-                if self.max_plate_nesting == float('inf'):
-                    raise ValueError("Automatic vectorization over num_particles requires " +
-                                     "a finite value for `max_plate_nesting` arg.")
-                self.max_plate_nesting += 1
+        if self.vectorize_particles and self.num_particles > 1:
+            self.max_plate_nesting += 1
         self.strict_enumeration_warning = strict_enumeration_warning
 
+    def _guess_max_plate_nesting(self, model, guide, *args, **kwargs):
+        """
+        Guesses max_plate_nesting by running the (model,guide) pair once
+        without enumeration. This optimistically assumes static model
+        structure.
+        """
+        # Ignore validation to allow model-enumerated sites absent from the guide.
+        with poutine.block():
+            guide_trace = poutine.trace(guide).get_trace(*args, **kwargs)
+            model_trace = poutine.trace(
+                poutine.replay(model, trace=guide_trace)).get_trace(*args, **kwargs)
+        guide_trace = prune_subsample_sites(guide_trace)
+        model_trace = prune_subsample_sites(model_trace)
+        sites = [site
+                 for trace in (model_trace, guide_trace)
+                 for site in trace.nodes.values()
+                 if site["type"] == "sample"]
+
+        # Validate shapes now, since shape constraints will be weaker once
+        # max_plate_nesting is changed from float('inf') to some finite value.
+        # Here we know the traces are not enumerated, but later we'll need to
+        # allow broadcasting of dims to the left of max_plate_nesting.
+        if is_validation_enabled():
+            guide_trace.compute_log_prob()
+            model_trace.compute_log_prob()
+            for site in sites:
+                check_site_shape(site, max_plate_nesting=float('inf'))
+
+        dims = [frame.dim
+                for site in sites
+                for frame in site["cond_indep_stack"]
+                if frame.vectorized]
+        self.max_plate_nesting = -min(dims) if dims else 0
+        if self.vectorize_particles and self.num_particles > 1:
+            self.max_plate_nesting += 1
+        logging.info('Guessed max_plate_nesting = {}'.format(self.max_plate_nesting))
+
     def _vectorized_num_particles(self, fn):
         """
         Wraps a callable inside an outermost :class:`~pyro.plate` to parallelize
@@ -113,6 +152,8 @@ def _get_traces(self, model, guide, *args, **kwargs):
         the result packaged as a trace generator.
         """
         if self.vectorize_particles:
+            if self.max_plate_nesting == float('inf'):
+                self._guess_max_plate_nesting(model, guide, *args, **kwargs)
             yield self._get_vectorized_trace(model, guide, *args, **kwargs)
         else:
             for i in range(self.num_particles):
diff --git a/pyro/infer/traceenum_elbo.py b/pyro/infer/traceenum_elbo.py
--- a/pyro/infer/traceenum_elbo.py
+++ b/pyro/infer/traceenum_elbo.py
@@ -15,10 +15,10 @@
 import pyro.poutine as poutine
 from pyro.distributions.torch_distribution import ReshapedDistribution
 from pyro.distributions.util import is_identically_zero, scale_and_mask
-from pyro.ops.contract import contract_tensor_tree, contract_to_tensor
 from pyro.infer.elbo import ELBO
 from pyro.infer.enum import get_importance_trace, iter_discrete_escape, iter_discrete_extend
 from pyro.infer.util import Dice, is_validation_enabled
+from pyro.ops.contract import contract_tensor_tree, contract_to_tensor
 from pyro.poutine.enumerate_messenger import EnumerateMessenger
 from pyro.util import check_traceenum_requirements, warn_if_nan
 
@@ -256,6 +256,8 @@ def _get_traces(self, model, guide, *args, **kwargs):
         Runs the guide and runs the model against the guide with
         the result packaged as a trace generator.
         """
+        if self.max_plate_nesting == float('inf'):
+            self._guess_max_plate_nesting(model, guide, *args, **kwargs)
         if self.vectorize_particles:
             guide = self._vectorized_num_particles(guide)
             model = self._vectorized_num_particles(model)
diff --git a/pyro/util.py b/pyro/util.py
--- a/pyro/util.py
+++ b/pyro/util.py
@@ -201,9 +201,9 @@ def check_model_guide_match(model_trace, guide_trace, max_plate_nesting=float('i
 
             # Allow broadcasting outside of max_plate_nesting.
             if len(model_shape) > max_plate_nesting:
-                model_shape = model_shape[len(model_shape) - max_plate_nesting:]
+                model_shape = model_shape[len(model_shape) - max_plate_nesting - model_site["fn"].event_dim:]
             if len(guide_shape) > max_plate_nesting:
-                guide_shape = guide_shape[len(guide_shape) - max_plate_nesting:]
+                guide_shape = guide_shape[len(guide_shape) - max_plate_nesting - guide_site["fn"].event_dim:]
             if model_shape == guide_shape:
                 continue
             for model_size, guide_size in zip_longest(reversed(model_shape), reversed(guide_shape), fillvalue=1):


Test patch:
diff --git a/tests/contrib/oed/test_eig.py b/tests/contrib/oed/test_eig.py
--- a/tests/contrib/oed/test_eig.py
+++ b/tests/contrib/oed/test_eig.py
@@ -20,7 +20,7 @@
     vi_ape, naive_rainforth_eig, donsker_varadhan_eig, barber_agakov_ape
 )
 from pyro.contrib.oed.util import linear_model_ground_truth
-from pyro.infer import TraceEnum_ELBO
+from pyro.infer import Trace_ELBO
 
 logger = logging.getLogger(__name__)
 
@@ -51,7 +51,7 @@
 # Aux
 ########################################################################################
 
-elbo = TraceEnum_ELBO(strict_enumeration_warning=False).differentiable_loss
+elbo = Trace_ELBO().differentiable_loss
 
 
 def bernoulli_guide(design):
diff --git a/tests/infer/test_enum.py b/tests/infer/test_enum.py
--- a/tests/infer/test_enum.py
+++ b/tests/infer/test_enum.py
@@ -133,8 +133,7 @@ def guide():
         with pyro.plate("batch", 3):
             pyro.sample("z", UnsafeBernoulli(p))
 
-    elbo = TraceEnum_ELBO(max_plate_nesting=1,
-                          strict_enumeration_warning=any([enumerate1]))
+    elbo = TraceEnum_ELBO(strict_enumeration_warning=any([enumerate1]))
     loss = elbo.loss(model, guide)
     assert not math.isnan(loss), loss
     loss = elbo.differentiable_loss(model, guide)
@@ -223,8 +222,7 @@ def test_svi_step_smoke(model, guide, enumerate1):
 
     guide = config_enumerate(guide, default=enumerate1)
     optimizer = pyro.optim.Adam({"lr": .001})
-    elbo = TraceEnum_ELBO(max_plate_nesting=1,
-                          strict_enumeration_warning=any([enumerate1]))
+    elbo = TraceEnum_ELBO(strict_enumeration_warning=any([enumerate1]))
     inference = SVI(model, guide, optimizer, loss=elbo)
     inference.step(data)
 
@@ -285,8 +283,7 @@ def guide():
         pyro.sample("loc", dist.Normal(loc, var))
         pyro.sample("b", dist.Bernoulli(p))
 
-    elbo = TraceEnum_ELBO(max_plate_nesting=1,
-                          strict_enumeration_warning=any([enumerate1]))
+    elbo = TraceEnum_ELBO(strict_enumeration_warning=any([enumerate1]))
     inference = SVI(model, guide, pyro.optim.Adam({}), elbo)
     inference.step()
 
@@ -313,8 +310,7 @@ def guide():
         with pyro.plate("particles", num_particles):
             pyro.sample("z", dist.Bernoulli(q).expand_by([num_particles]))
 
-    elbo = TraceEnum_ELBO(max_plate_nesting=1,
-                          strict_enumeration_warning=any([enumerate1]))
+    elbo = TraceEnum_ELBO(strict_enumeration_warning=any([enumerate1]))
 
     if method == "loss":
         actual = elbo.loss(model, guide) / num_particles
@@ -356,8 +352,7 @@ def guide():
         with pyro.plate("particles", num_particles):
             pyro.sample("z", dist.Normal(q, 1.).expand_by([num_particles]))
 
-    elbo = TraceEnum_ELBO(max_plate_nesting=1,
-                          strict_enumeration_warning=any([enumerate1]))
+    elbo = TraceEnum_ELBO(strict_enumeration_warning=any([enumerate1]))
 
     if method == "loss":
         actual = elbo.loss(model, guide) / num_particles
@@ -420,8 +415,7 @@ def guide():
     expected_loss = kl.item()
     expected_grad = grad(kl, [q])[0]
 
-    elbo = TraceEnum_ELBO(max_plate_nesting=0,
-                          num_particles=num_particles,
+    elbo = TraceEnum_ELBO(num_particles=num_particles,
                           vectorize_particles=True,
                           strict_enumeration_warning=any([enumerate1, enumerate2]))
     if method == "differentiable_loss":
@@ -467,8 +461,7 @@ def guide():
     expected_loss = kl.item()
     expected_grad = grad(kl, [q])[0]
 
-    elbo = TraceEnum_ELBO(max_plate_nesting=0,
-                          num_particles=num_particles,
+    elbo = TraceEnum_ELBO(num_particles=num_particles,
                           vectorize_particles=True,
                           strict_enumeration_warning=any([enumerate1, enumerate2, enumerate3]))
     if method == "differentiable_loss":
@@ -559,8 +552,7 @@ def guide():
     expected_loss = kl.item()
     expected_grad = grad(kl, [q])[0]
 
-    elbo = TraceEnum_ELBO(max_plate_nesting=0,
-                          num_particles=num_particles,
+    elbo = TraceEnum_ELBO(num_particles=num_particles,
                           vectorize_particles=True,
                           strict_enumeration_warning=any([enumerate1, enumerate2, enumerate3]))
     if method == "differentiable_loss":
@@ -609,8 +601,7 @@ def guide():
     expected_loss = kl.item()
     expected_grad = grad(kl, [q])[0]
 
-    elbo = TraceEnum_ELBO(max_plate_nesting=2,
-                          strict_enumeration_warning=any([enumerate1, enumerate2]))
+    elbo = TraceEnum_ELBO(strict_enumeration_warning=any([enumerate1, enumerate2]))
     actual_loss = elbo.loss_and_grads(model, guide) / num_particles
     actual_grad = pyro.param('q').grad / num_particles
 
@@ -629,7 +620,7 @@ def guide():
 @pytest.mark.parametrize("irange_dim", [1, 2])
 def test_elbo_irange(irange_dim, enumerate1, enumerate2):
     pyro.clear_param_store()
-    num_particles = 1 if all([enumerate1, enumerate2]) else 10000
+    num_particles = 1 if all([enumerate1, enumerate2]) else 20000
     q = pyro.param("q", torch.tensor(0.75, requires_grad=True))
     p = 0.2693204236205713  # for which kl(Bernoulli(q), Bernoulli(p)) = 0.5
 
@@ -652,8 +643,7 @@ def guide():
     expected_loss = kl.item()
     expected_grad = grad(kl, [q])[0]
 
-    elbo = TraceEnum_ELBO(max_plate_nesting=1,
-                          strict_enumeration_warning=any([enumerate1, enumerate2]))
+    elbo = TraceEnum_ELBO(strict_enumeration_warning=any([enumerate1, enumerate2]))
     actual_loss = elbo.loss_and_grads(model, guide) / num_particles
     actual_grad = pyro.param('q').grad / num_particles
 
@@ -708,8 +698,7 @@ def guide():
     expected_loss = kl.item()
     expected_grad = grad(kl, [q])[0]
 
-    elbo = TraceEnum_ELBO(max_plate_nesting=3,
-                          num_particles=num_particles,
+    elbo = TraceEnum_ELBO(num_particles=num_particles,
                           vectorize_particles=True,
                           strict_enumeration_warning=any([enumerate1, enumerate2, enumerate3]))
     actual_loss = elbo.loss_and_grads(model, guide)
@@ -760,8 +749,7 @@ def guide():
     expected_loss = kl.item()
     expected_grad = grad(kl, [q])[0]
 
-    elbo = TraceEnum_ELBO(max_plate_nesting=2,
-                          strict_enumeration_warning=any([enumerate1, enumerate2, enumerate3]))
+    elbo = TraceEnum_ELBO(strict_enumeration_warning=any([enumerate1, enumerate2, enumerate3]))
     actual_loss = elbo.loss_and_grads(model, guide) / num_particles
     actual_grad = pyro.param('q').grad / num_particles
 
@@ -812,8 +800,7 @@ def guide():
     expected_loss = kl.item()
     expected_grad = grad(kl, [q])[0]
 
-    elbo = TraceEnum_ELBO(max_plate_nesting=2,
-                          strict_enumeration_warning=any([enumerate1, enumerate2, enumerate3]))
+    elbo = TraceEnum_ELBO(strict_enumeration_warning=any([enumerate1, enumerate2, enumerate3]))
     actual_loss = elbo.loss_and_grads(model, guide) / num_particles
     actual_grad = pyro.param('q').grad / num_particles
 
@@ -864,8 +851,7 @@ def guide():
     expected_loss = kl.item()
     expected_grad = grad(kl, [q])[0]
 
-    elbo = TraceEnum_ELBO(max_plate_nesting=1,
-                          strict_enumeration_warning=any([enumerate1, enumerate2, enumerate3]))
+    elbo = TraceEnum_ELBO(strict_enumeration_warning=any([enumerate1, enumerate2, enumerate3]))
     actual_loss = elbo.loss_and_grads(model, guide) / num_particles
     actual_grad = pyro.param('q').grad / num_particles
 
@@ -899,8 +885,7 @@ def guide():
             pyro.sample("z", dist.Bernoulli(q2 * y + 0.10))
 
     logger.info("Computing gradients using surrogate loss")
-    elbo = TraceEnum_ELBO(max_plate_nesting=1,
-                          strict_enumeration_warning=any([enumerate1]))
+    elbo = TraceEnum_ELBO(strict_enumeration_warning=any([enumerate1]))
     elbo.loss_and_grads(model, config_enumerate(guide, default=enumerate1))
     actual_grad_q1 = pyro.param('q1').grad / num_particles
     actual_grad_q2 = pyro.param('q2').grad / num_particles
@@ -955,8 +940,7 @@ def guide():
                 pyro.sample("z", dist.Normal(q2 * y + 0.10, 1.0))
 
     logger.info("Computing gradients using surrogate loss")
-    elbo = TraceEnum_ELBO(max_plate_nesting=1,
-                          strict_enumeration_warning=any([enumerate1]))
+    elbo = TraceEnum_ELBO(strict_enumeration_warning=any([enumerate1]))
     elbo.loss_and_grads(model, guide)
     actual_grad_q1 = pyro.param('q1').grad / num_particles
     if include_z:
@@ -1018,8 +1002,7 @@ def guide(num_particles):
 
     for ed, num_particles in zip([None, 'parallel', 'sequential'], [30000, 20000, 20000]):
         pyro.clear_param_store()
-        elbo = TraceEnum_ELBO(max_plate_nesting=1,
-                              strict_enumeration_warning=any([ed]))
+        elbo = TraceEnum_ELBO(strict_enumeration_warning=any([ed]))
         elbo.loss_and_grads(model, config_enumerate(guide, default=ed), num_particles)
         results[str(ed)] = {}
         for q in qs:
@@ -1060,8 +1043,7 @@ def guide():
             pyro.sample("z", dist.Bernoulli(q).expand_by([num_particles]))
             pyro.sample("y", ShapeAugmentedGamma(a, torch.tensor(1.0)).expand_by([num_particles]))
 
-    elbo = TraceEnum_ELBO(max_plate_nesting=1,
-                          strict_enumeration_warning=any([enumerate1]))
+    elbo = TraceEnum_ELBO(strict_enumeration_warning=any([enumerate1]))
     elbo.loss_and_grads(model, guide)
 
     actual_q = q.grad / num_particles
@@ -1117,7 +1099,7 @@ def guide(data):
         for i in range(num_steps):
             pyro.sample("x_{}".format(i), dist.Categorical(mean_field_probs[i]))
 
-    elbo = TraceEnum_ELBO(max_plate_nesting=0)
+    elbo = TraceEnum_ELBO()
     elbo.loss_and_grads(model, guide, data)
 
     expected_unconstrained_grads = {
@@ -1179,7 +1161,7 @@ def guide(data):
             probs = init_probs if x is None else transition_probs[x]
             x = pyro.sample("x_{}".format(i), dist.Categorical(probs))
 
-    elbo = TraceEnum_ELBO(max_plate_nesting=0)
+    elbo = TraceEnum_ELBO()
     elbo.loss_and_grads(model, guide, data)
 
     # These golden values simply test agreement between parallel and sequential.
@@ -1242,7 +1224,7 @@ def model(data):
     def guide(data):
         pass
 
-    elbo = TraceEnum_ELBO(max_plate_nesting=0)
+    elbo = TraceEnum_ELBO()
     elbo.differentiable_loss(model, guide, data)
 
 
@@ -1272,7 +1254,7 @@ def guide(data):
         pyro.sample("x", dist.Categorical(init_probs),
                     infer={"enumerate": "parallel"})
 
-    elbo = TraceEnum_ELBO(max_plate_nesting=0)
+    elbo = TraceEnum_ELBO()
     elbo.differentiable_loss(model, guide, data)
 
 
@@ -1332,7 +1314,7 @@ def guide():
         probs_x = pyro.param("guide_probs_x")
         pyro.sample("x", dist.Categorical(probs_x))
 
-    elbo = TraceEnum_ELBO(max_plate_nesting=0, strict_enumeration_warning=False)
+    elbo = TraceEnum_ELBO(strict_enumeration_warning=False)
     auto_loss = elbo.differentiable_loss(auto_model, guide)
     hand_loss = elbo.differentiable_loss(hand_model, guide)
     _check_loss_and_grads(hand_loss, auto_loss)
@@ -1378,7 +1360,7 @@ def guide():
         probs_x = pyro.param("guide_probs_x")
         pyro.sample("x", dist.Categorical(probs_x))
 
-    elbo = TraceEnum_ELBO(max_plate_nesting=0, strict_enumeration_warning=False)
+    elbo = TraceEnum_ELBO(strict_enumeration_warning=False)
     auto_loss = elbo.differentiable_loss(auto_model, guide)
     hand_loss = elbo.differentiable_loss(hand_model, guide)
     _check_loss_and_grads(hand_loss, auto_loss)
@@ -1423,7 +1405,7 @@ def guide():
         probs_x = pyro.param("guide_probs_x")
         pyro.sample("x", dist.Categorical(probs_x))
 
-    elbo = TraceEnum_ELBO(max_plate_nesting=0, strict_enumeration_warning=False)
+    elbo = TraceEnum_ELBO(strict_enumeration_warning=False)
     auto_loss = elbo.differentiable_loss(auto_model, guide)
     hand_loss = elbo.differentiable_loss(hand_model, guide)
     _check_loss_and_grads(hand_loss, auto_loss)
diff --git a/tests/infer/test_gradient.py b/tests/infer/test_gradient.py
--- a/tests/infer/test_gradient.py
+++ b/tests/infer/test_gradient.py
@@ -63,7 +63,7 @@ def guide(subsample):
         num_particles = 1
 
     optim = Adam({"lr": 0.1})
-    elbo = Elbo(max_plate_nesting=1,
+    elbo = Elbo(max_plate_nesting=1,  # set this to ensure rng agrees across runs
                 num_particles=num_particles,
                 vectorize_particles=True,
                 strict_enumeration_warning=False)
@@ -89,7 +89,7 @@ def guide(subsample):
 def test_plate(Elbo, reparameterized):
     pyro.clear_param_store()
     data = torch.tensor([-0.5, 2.0])
-    num_particles = 100000
+    num_particles = 200000
     precision = 0.06
     Normal = dist.Normal if reparameterized else fakes.NonreparameterizedNormal
 
@@ -117,7 +117,8 @@ def guide():
         pyro.sample("nuisance_a", Normal(0, 1))
 
     optim = Adam({"lr": 0.1})
-    inference = SVI(model, guide, optim, loss=Elbo(strict_enumeration_warning=False))
+    elbo = Elbo(strict_enumeration_warning=False)
+    inference = SVI(model, guide, optim, loss=elbo)
     inference.loss_and_grads(model, guide)
     params = dict(pyro.get_param_store().named_parameters())
     actual_grads = {name: param.grad.detach().cpu().numpy() / num_particles
@@ -161,8 +162,7 @@ def guide():
         pyro.sample("nuisance_a", Normal(0, 1))
 
     optim = Adam({"lr": 0.1})
-    loss = Elbo(max_plate_nesting=1,
-                num_particles=num_particles,
+    loss = Elbo(num_particles=num_particles,
                 vectorize_particles=True,
                 strict_enumeration_warning=False)
     inference = SVI(model, guide, optim, loss=loss)
diff --git a/tests/infer/test_valid_models.py b/tests/infer/test_valid_models.py
--- a/tests/infer/test_valid_models.py
+++ b/tests/infer/test_valid_models.py
@@ -108,14 +108,15 @@ def test_model_guide_dim_mismatch_error(Elbo):
     def model():
         loc = torch.zeros(2)
         scale = torch.ones(2)
-        pyro.sample("x", dist.Normal(loc, scale))
+        pyro.sample("x", dist.Normal(loc, scale).independent(1))
 
     def guide():
         loc = pyro.param("loc", torch.zeros(2, 1, requires_grad=True))
         scale = pyro.param("scale", torch.ones(2, 1, requires_grad=True))
-        pyro.sample("x", dist.Normal(loc, scale))
+        pyro.sample("x", dist.Normal(loc, scale).independent(2))
 
-    assert_error(model, guide, Elbo(), match='Model and guide shapes disagree')
+    assert_error(model, guide, Elbo(strict_enumeration_warning=False),
+                 match='invalid log_prob shape|Model and guide event_dims disagree')
 
 
 @pytest.mark.parametrize("Elbo", [Trace_ELBO, TraceGraph_ELBO, TraceEnum_ELBO])
@@ -124,14 +125,15 @@ def test_model_guide_shape_mismatch_error(Elbo):
     def model():
         loc = torch.zeros(1, 2)
         scale = torch.ones(1, 2)
-        pyro.sample("x", dist.Normal(loc, scale))
+        pyro.sample("x", dist.Normal(loc, scale).independent(2))
 
     def guide():
         loc = pyro.param("loc", torch.zeros(2, 1, requires_grad=True))
         scale = pyro.param("scale", torch.ones(2, 1, requires_grad=True))
-        pyro.sample("x", dist.Normal(loc, scale))
+        pyro.sample("x", dist.Normal(loc, scale).independent(2))
 
-    assert_error(model, guide, Elbo(), match='Model and guide shapes disagree')
+    assert_error(model, guide, Elbo(strict_enumeration_warning=False),
+                 match='Model and guide shapes disagree')
 
 
 @pytest.mark.parametrize("Elbo", [Trace_ELBO, TraceGraph_ELBO, TraceEnum_ELBO])
@@ -229,9 +231,10 @@ def guide():
     assert_ok(model, guide, Elbo())
 
 
+@pytest.mark.parametrize("max_plate_nesting", [0, float('inf')])
 @pytest.mark.parametrize("subsample_size", [None, 2], ids=["full", "subsample"])
 @pytest.mark.parametrize("Elbo", [Trace_ELBO, TraceGraph_ELBO, TraceEnum_ELBO])
-def test_irange_irange_ok(subsample_size, Elbo):
+def test_irange_irange_ok(subsample_size, Elbo, max_plate_nesting):
 
     def model():
         p = torch.tensor(0.5)
@@ -252,12 +255,13 @@ def guide():
     if Elbo is TraceEnum_ELBO:
         guide = config_enumerate(guide, "parallel")
 
-    assert_ok(model, guide, Elbo(max_plate_nesting=0))
+    assert_ok(model, guide, Elbo(max_plate_nesting=max_plate_nesting))
 
 
+@pytest.mark.parametrize("max_plate_nesting", [0, float('inf')])
 @pytest.mark.parametrize("subsample_size", [None, 2], ids=["full", "subsample"])
 @pytest.mark.parametrize("Elbo", [Trace_ELBO, TraceGraph_ELBO, TraceEnum_ELBO])
-def test_irange_irange_swap_ok(subsample_size, Elbo):
+def test_irange_irange_swap_ok(subsample_size, Elbo, max_plate_nesting):
 
     def model():
         p = torch.tensor(0.5)
@@ -278,7 +282,7 @@ def guide():
     if Elbo is TraceEnum_ELBO:
         guide = config_enumerate(guide, "parallel")
 
-    assert_ok(model, guide, Elbo(max_plate_nesting=0))
+    assert_ok(model, guide, Elbo(max_plate_nesting=max_plate_nesting))
 
 
 @pytest.mark.parametrize("subsample_size", [None, 5], ids=["full", "subsample"])
@@ -453,7 +457,7 @@ def model():
                 pyro.sample("z", dist.Bernoulli(p).expand_by([len(ind_inner), 1]))  # error here
 
     guide = config_enumerate(model) if Elbo is TraceEnum_ELBO else model
-    assert_error(model, guide, Elbo(), match='invalid log_prob shape')
+    assert_error(model, guide, Elbo(), match='invalid log_prob shape|shape mismatch')
 
 
 @pytest.mark.parametrize("Elbo", [Trace_ELBO, TraceGraph_ELBO, TraceEnum_ELBO])
@@ -666,35 +670,40 @@ def guide():
                  match='invalid log_prob shape')
 
 
-@pytest.mark.parametrize('max_plate_nesting', [0, 1, 2])
+@pytest.mark.parametrize('max_plate_nesting', [0, 1, 2, float('inf')])
 def test_enum_discrete_parallel_ok(max_plate_nesting):
-    plate_shape = torch.Size([1] * max_plate_nesting)
+    guessed_nesting = 0 if max_plate_nesting == float('inf') else max_plate_nesting
+    plate_shape = torch.Size([1] * guessed_nesting)
 
     def model():
         p = torch.tensor(0.5)
         x = pyro.sample("x", dist.Bernoulli(p))
-        assert x.shape == torch.Size([2]) + plate_shape
+        if max_plate_nesting != float('inf'):
+            assert x.shape == torch.Size([2]) + plate_shape
 
     def guide():
         p = pyro.param("p", torch.tensor(0.5, requires_grad=True))
         x = pyro.sample("x", dist.Bernoulli(p))
-        assert x.shape == torch.Size([2]) + plate_shape
+        if max_plate_nesting != float('inf'):
+            assert x.shape == torch.Size([2]) + plate_shape
 
     assert_ok(model, config_enumerate(guide, "parallel"),
               TraceEnum_ELBO(max_plate_nesting=max_plate_nesting))
 
 
-@pytest.mark.parametrize('max_plate_nesting', [0, 1, 2])
+@pytest.mark.parametrize('max_plate_nesting', [0, 1, 2, float('inf')])
 def test_enum_discrete_parallel_nested_ok(max_plate_nesting):
-    plate_shape = torch.Size([1] * max_plate_nesting)
+    guessed_nesting = 0 if max_plate_nesting == float('inf') else max_plate_nesting
+    plate_shape = torch.Size([1] * guessed_nesting)
 
     def model():
         p2 = torch.ones(2) / 2
         p3 = torch.ones(3) / 3
         x2 = pyro.sample("x2", dist.OneHotCategorical(p2))
         x3 = pyro.sample("x3", dist.OneHotCategorical(p3))
-        assert x2.shape == torch.Size([2]) + plate_shape + p2.shape
-        assert x3.shape == torch.Size([3, 1]) + plate_shape + p3.shape
+        if max_plate_nesting != float('inf'):
+            assert x2.shape == torch.Size([2]) + plate_shape + p2.shape
+            assert x3.shape == torch.Size([3, 1]) + plate_shape + p3.shape
 
     assert_ok(model, config_enumerate(model, "parallel"),
               TraceEnum_ELBO(max_plate_nesting=max_plate_nesting))
@@ -761,9 +770,10 @@ def model():
     assert_ok(model, guide, elbo)
 
 
+@pytest.mark.parametrize('max_plate_nesting', [1, float('inf')])
 @pytest.mark.parametrize('enumerate_', [None, "sequential", "parallel"])
 @pytest.mark.parametrize('is_validate', [True, False])
-def test_enum_discrete_plate_dependency_warning(enumerate_, is_validate):
+def test_enum_discrete_plate_dependency_warning(enumerate_, is_validate, max_plate_nesting):
 
     def model():
         pyro.sample("w", dist.Bernoulli(0.5), infer={'enumerate': 'parallel'})
@@ -773,14 +783,16 @@ def model():
         pyro.sample("y", dist.Bernoulli(x.mean()))  # user should move this line up
 
     with pyro.validation_enabled(is_validate):
+        elbo = TraceEnum_ELBO(max_plate_nesting=max_plate_nesting)
         if enumerate_ and is_validate:
-            assert_warning(model, model, TraceEnum_ELBO(max_plate_nesting=1))
+            assert_warning(model, model, elbo)
         else:
-            assert_ok(model, model, TraceEnum_ELBO(max_plate_nesting=1))
+            assert_ok(model, model, elbo)
 
 
+@pytest.mark.parametrize('max_plate_nesting', [1, float('inf')])
 @pytest.mark.parametrize('enumerate_', [None, "sequential", "parallel"])
-def test_enum_discrete_irange_plate_dependency_ok(enumerate_):
+def test_enum_discrete_irange_plate_dependency_ok(enumerate_, max_plate_nesting):
 
     def model():
         pyro.sample("w", dist.Bernoulli(0.5), infer={'enumerate': 'parallel'})
@@ -791,12 +803,13 @@ def model():
                 pyro.sample("x_{}".format(i), dist.Bernoulli(0.5).expand_by([5]),
                             infer={'enumerate': enumerate_})
 
-    assert_ok(model, model, TraceEnum_ELBO(max_plate_nesting=1))
+    assert_ok(model, model, TraceEnum_ELBO(max_plate_nesting=max_plate_nesting))
 
 
+@pytest.mark.parametrize('max_plate_nesting', [1, float('inf')])
 @pytest.mark.parametrize('enumerate_', [None, "sequential", "parallel"])
 @pytest.mark.parametrize('is_validate', [True, False])
-def test_enum_discrete_iranges_plate_dependency_warning(enumerate_, is_validate):
+def test_enum_discrete_iranges_plate_dependency_warning(enumerate_, is_validate, max_plate_nesting):
 
     def model():
         pyro.sample("w", dist.Bernoulli(0.5), infer={'enumerate': 'parallel'})
@@ -811,10 +824,11 @@ def model():
             pyro.sample("y_{}".format(i), dist.Bernoulli(0.5))
 
     with pyro.validation_enabled(is_validate):
+        elbo = TraceEnum_ELBO(max_plate_nesting=max_plate_nesting)
         if enumerate_ and is_validate:
-            assert_warning(model, model, TraceEnum_ELBO(max_plate_nesting=1))
+            assert_warning(model, model, elbo)
         else:
-            assert_ok(model, model, TraceEnum_ELBO(max_plate_nesting=1))
+            assert_ok(model, model, elbo)
 
 
 @pytest.mark.parametrize('enumerate_', [None, "sequential", "parallel"])
@@ -966,7 +980,7 @@ def model():
             assert q.shape == (8, 5, 7, 6)
 
     guide = config_enumerate(model, expand=expand) if enumerate_ else model
-    assert_ok(model, guide, Elbo())
+    assert_ok(model, guide, Elbo(max_plate_nesting=4))
 
 
 @pytest.mark.parametrize("Elbo,expand", [