Repository: pyro-ppl/numpyro

Issue:
AutoDelta raises duplication error
Currently, using AutoDelta on some models generates sort of `AssertionError: all sites must have unique names but got `sd_dip_rnde` duplicated` error. More information to come...


Discussion:

Resolving patch:
diff --git a/examples/neutra.py b/examples/neutra.py
--- a/examples/neutra.py
+++ b/examples/neutra.py
@@ -116,28 +116,28 @@ def main(args):
     ax1.set_title('Autoguide training loss\n(after 1000 steps)')
 
     ax2.contourf(X1, X2, P, cmap='OrRd')
-    sns.kdeplot(guide_samples[:, 0], guide_samples[:, 1], n_levels=30, ax=ax2)
+    sns.kdeplot(x=guide_samples[:, 0], y=guide_samples[:, 1], n_levels=30, ax=ax2)
     ax2.set(xlim=[-3, 3], ylim=[-3, 3],
             xlabel='x0', ylabel='x1', title='Posterior using\nAutoBNAFNormal guide')
 
-    sns.scatterplot(guide_base_samples[:, 0], guide_base_samples[:, 1], ax=ax3,
+    sns.scatterplot(x=guide_base_samples[:, 0], y=guide_base_samples[:, 1], ax=ax3,
                     hue=guide_trans_samples[:, 0] < 0.)
     ax3.set(xlim=[-3, 3], ylim=[-3, 3],
             xlabel='x0', ylabel='x1', title='AutoBNAFNormal base samples\n(True=left moon; False=right moon)')
 
     ax4.contourf(X1, X2, P, cmap='OrRd')
-    sns.kdeplot(vanilla_samples[:, 0], vanilla_samples[:, 1], n_levels=30, ax=ax4)
+    sns.kdeplot(x=vanilla_samples[:, 0], y=vanilla_samples[:, 1], n_levels=30, ax=ax4)
     ax4.plot(vanilla_samples[-50:, 0], vanilla_samples[-50:, 1], 'bo-', alpha=0.5)
     ax4.set(xlim=[-3, 3], ylim=[-3, 3],
             xlabel='x0', ylabel='x1', title='Posterior using\nvanilla HMC sampler')
 
-    sns.scatterplot(zs[:, 0], zs[:, 1], ax=ax5, hue=samples[:, 0] < 0.,
+    sns.scatterplot(x=zs[:, 0], y=zs[:, 1], ax=ax5, hue=samples[:, 0] < 0.,
                     s=30, alpha=0.5, edgecolor="none")
     ax5.set(xlim=[-5, 5], ylim=[-5, 5],
             xlabel='x0', ylabel='x1', title='Samples from the\nwarped posterior - p(z)')
 
     ax6.contourf(X1, X2, P, cmap='OrRd')
-    sns.kdeplot(samples[:, 0], samples[:, 1], n_levels=30, ax=ax6)
+    sns.kdeplot(x=samples[:, 0], y=samples[:, 1], n_levels=30, ax=ax6)
     ax6.plot(samples[-50:, 0], samples[-50:, 1], 'bo-', alpha=0.2)
     ax6.set(xlim=[-3, 3], ylim=[-3, 3],
             xlabel='x0', ylabel='x1', title='Posterior using\nNeuTra HMC sampler')
diff --git a/numpyro/infer/autoguide.py b/numpyro/infer/autoguide.py
--- a/numpyro/infer/autoguide.py
+++ b/numpyro/infer/autoguide.py
@@ -287,15 +287,15 @@ class AutoDelta(AutoGuide):
     """
     def __init__(self, model, *, prefix='auto', init_loc_fn=init_to_median,
                  create_plates=None):
-        self.init_loc_fn = init_loc_fn
         self._event_dims = {}
         super().__init__(model, prefix=prefix, init_loc_fn=init_loc_fn, create_plates=create_plates)
 
     def _setup_prototype(self, *args, **kwargs):
         super()._setup_prototype(*args, **kwargs)
-        self._init_locs = {
-            k: v for k, v in self._postprocess_fn(self._init_locs).items() if k in self._init_locs
-        }
+        with numpyro.handlers.block():
+            self._init_locs = {
+                k: v for k, v in self._postprocess_fn(self._init_locs).items() if k in self._init_locs
+            }
         for name, site in self.prototype_trace.items():
             if site["type"] != "sample" or site["is_observed"]:
                 continue


Test patch:
diff --git a/test/infer/test_autoguide.py b/test/infer/test_autoguide.py
--- a/test/infer/test_autoguide.py
+++ b/test/infer/test_autoguide.py
@@ -92,7 +92,7 @@ def test_logistic_regression(auto_class, Elbo):
 
     def model(data, labels):
         coefs = numpyro.sample('coefs', dist.Normal(jnp.zeros(dim), jnp.ones(dim)))
-        logits = jnp.sum(coefs * data, axis=-1)
+        logits = numpyro.deterministic("logits", jnp.sum(coefs * data, axis=-1))
         return numpyro.sample('obs', dist.Bernoulli(logits=logits), obs=labels)
 
     adam = optim.Adam(0.01)
diff --git a/test/pyroapi/test_pyroapi.py b/test/pyroapi/test_pyroapi.py
--- a/test/pyroapi/test_pyroapi.py
+++ b/test/pyroapi/test_pyroapi.py
@@ -8,7 +8,7 @@
 pytestmark = pytest.mark.filterwarnings("ignore::numpyro.compat.util.UnsupportedAPIWarning")
 
 
-@pytest.yield_fixture
+@pytest.fixture
 def backend():
     with pyro_backend('numpy'):
         yield