Repository: pyro-ppl/numpyro

Issue:
Error using AutoDiagonalNormal with SVI + Predictive
I'm getting an error when using `AutoDiagonalNormal` to construct a guide for SVI, and then subsequently using that guide to construct the predictive distribution. 

```
UnboundLocalError: local variable 'event_ndim' referenced before assignment
```

The error comes from [this section](https://github.com/pyro-ppl/numpyro/blob/6823c495f62774172936fb437bb17210fd0179e0/numpyro/infer/autoguide.py#L485-L497) of the code. It looks to me like if the `if` statement is satisfied, then lines 495-497 will always fail? So it doesn't seem quite right, but I don't understand what's happening there well enough yet to be sure how to fix it.

Here's a minimal example. I see the error with NumPyro 0.6.0 and also when installing from `master`. Swapping `AutoDiagonalNormal` for `AutoNormal` resolves the issue, even though for this example where there is only a single latent I would expect them to be equivalent?

```python
import numpyro
import numpyro.distributions as dist
from jax import random
from numpyro.infer import SVI, Trace_ELBO
from numpyro.infer.autoguide import AutoDiagonalNormal, AutoNormal
from numpyro.infer.util import Predictive


def model(x=None):
    n = x.shape[0] if x is not None else 1
    mu = numpyro.sample("mu", dist.Cauchy())

    with numpyro.plate("plate", n):
        numpyro.sample("x", dist.Normal(mu, 1), obs=x)


if __name__ == "__main__":
    rng_key = random.PRNGKey(0)
    alpha = 2

    rng_key, subkey = random.split(rng_key)
    x = random.normal(subkey, (1000,)) + alpha

    # AutoDiagonalNormal results in an error when sampling from the predictive
    # AutoNormal doesn't have the same problem
    # guide = AutoNormal(model)
    guide = AutoDiagonalNormal(model)
    svi = SVI(model, guide, numpyro.optim.Adam(1e-3), Trace_ELBO())

    rng_key, subkey = random.split(rng_key)
    params, losses = svi.run(subkey, 1000, x)

    predictive = Predictive(
        model, guide=guide, params=params, num_samples=1000
    )

    rng_key, subkey = random.split(rng_key)
    # this line results in error is AutoDiagonalNormal was used
    posterior_predictive_samples = predictive(subkey)
```


Discussion:
Thanks, @tcbegley! That's a bug introduced recently. Could you submit a PR with `event_ndim` moved outside of that `if/else` block?
```
event_ndim = site["fn"].event_dim
if ...
else ...
```

Resolving patch:
diff --git a/numpyro/infer/autoguide.py b/numpyro/infer/autoguide.py
--- a/numpyro/infer/autoguide.py
+++ b/numpyro/infer/autoguide.py
@@ -482,13 +482,13 @@ def __call__(self, *args, **kwargs):
             site = self.prototype_trace[name]
             transform = biject_to(site["fn"].support)
             value = transform(unconstrained_value)
+            event_ndim = site["fn"].event_dim
             if numpyro.get_mask() is False:
                 log_density = 0.0
             else:
                 log_density = -transform.log_abs_det_jacobian(
                     unconstrained_value, value
                 )
-                event_ndim = site["fn"].event_dim
                 log_density = sum_rightmost(
                     log_density, jnp.ndim(log_density) - jnp.ndim(value) + event_ndim
                 )


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
@@ -31,6 +31,7 @@
 )
 from numpyro.infer.initialization import init_to_median
 from numpyro.infer.reparam import TransformReparam
+from numpyro.infer.util import Predictive
 from numpyro.nn.auto_reg_nn import AutoregressiveNN
 from numpyro.util import fori_loop
 
@@ -75,6 +76,16 @@ def body_fn(i, val):
     )
     assert_allclose(jnp.mean(posterior_samples["beta"], 0), true_coefs, atol=0.05)
 
+    # Predictive can be instantiated from posterior samples...
+    predictive = Predictive(model, posterior_samples=posterior_samples)
+    predictive_samples = predictive(random.PRNGKey(1), None)
+    assert predictive_samples["obs"].shape == (1000, 2)
+
+    # ... or from the guide + params
+    predictive = Predictive(model, guide=guide, params=params, num_samples=1000)
+    predictive_samples = predictive(random.PRNGKey(1), None)
+    assert predictive_samples["obs"].shape == (1000, 2)
+
 
 @pytest.mark.parametrize(
     "auto_class",