Repository: pyro-ppl/numpyro

Issue:
Guide samples for AutoDelta don't contain deterministic variables 
I noticed that samples from the `AutoDelta` guide don't contain deterministic transforms of random variables in the model. Here's a minimal example which reproduces the issue:

```Python
import numpyro
import numpyro.distributions as dist
from numpyro import optim
from numpyro.infer import SVI
from numpyro.infer.autoguide import AutoDelta
from numpyro.infer.elbo import Trace_ELBO
from jax import random

def model():
    x = numpyro.sample('x', dist.Normal(0, 1).expand([10]))
    numpyro.deterministic('x2', x**2)
    
guide = AutoDelta(model)
svi = SVI(model, guide, optim.Adam(0.01), Trace_ELBO())
svi_state = svi.init(random.PRNGKey(0))
svi_result = svi.run(random.PRNGKey(1), 1000)
guide_samples = guide.sample_posterior(random.PRNGKey(2), svi_result.params)
print(guide_samples['x2'])
```
Is this related to #946?


Discussion:
Currently, `AutoDelta.sample_posterior` does not return deterministic sites. To collect those sites, we need to run `Predictive` on the `model` given those posterior samples.
Since this is marked as a good first issue, I could give this a try as a first contribution.

To run `Predictive` on the `model`, I would pass the model's `*args` and `**kwargs` to `AutoDelta.sample_posterior` as well. My first attempt is something like [this](https://gist.github.com/nikmich1/cb6dac89159b6bfd1343f13f3b1f89ca). If someone could take a brief look if this makes sense, before I create a pull request, this would be highly appreciated!
Hi @nikmich1, the solution looks good. To account for the non-trivial case of `sample_shape`, you can add `batch_ndims=len(sample_shape)` to the Predictive utility.
@fehiepsi Thank you very much for the feedback! I added the `sample_shape` as suggested. 

While running the unit tests, I noticed that `sample_posterior` should only pick up deterministic and not observed variables to be consistent e.g. with `AutoNormal`. Does it make sense to only return variables with `AutoDelta.prototype_trace[variable]["type"] == "deterministic"`  (something like [this](https://gist.github.com/nikmich1/cb6dac89159b6bfd1343f13f3b1f89ca))?
LGTM. You can obtain deterministic sites in advance so that we can
+ skip running Predictive if there are no deterministic sites
+ specify return_sites keywork in the Predictive constructor.

Resolving patch:
diff --git a/numpyro/infer/autoguide.py b/numpyro/infer/autoguide.py
--- a/numpyro/infer/autoguide.py
+++ b/numpyro/infer/autoguide.py
@@ -44,6 +44,7 @@
     periodic_repeat,
     sum_rightmost,
 )
+from numpyro.infer import Predictive
 from numpyro.infer.elbo import Trace_ELBO
 from numpyro.infer.initialization import init_to_median, init_to_uniform
 from numpyro.infer.util import helpful_support_errors, initialize_model
@@ -455,12 +456,25 @@ def __call__(self, *args, **kwargs):
 
         return result
 
-    def sample_posterior(self, rng_key, params, sample_shape=()):
+    def sample_posterior(self, rng_key, params, *args, sample_shape=(), **kwargs):
         locs = {k: params["{}_{}_loc".format(k, self.prefix)] for k in self._init_locs}
         latent_samples = {
             k: jnp.broadcast_to(v, sample_shape + jnp.shape(v)) for k, v in locs.items()
         }
-        return latent_samples
+        deterministic_vars = [
+            k for k, v in self.prototype_trace.items() if v["type"] == "deterministic"
+        ]
+        if not deterministic_vars:
+            return latent_samples
+        else:
+            predictive = Predictive(
+                model=self.model,
+                posterior_samples=latent_samples,
+                return_sites=deterministic_vars,
+                batch_ndims=len(sample_shape),
+            )
+            deterministic_samples = predictive(rng_key, *args, **kwargs)
+            return {**latent_samples, **deterministic_samples}
 
     def median(self, params):
         locs = {k: params["{}_{}_loc".format(k, self.prefix)] for k in self._init_locs}


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
@@ -141,11 +141,12 @@ def test_logistic_regression(auto_class, Elbo):
     logits = jnp.sum(true_coefs * data, axis=-1)
     labels = dist.Bernoulli(logits=logits).sample(random.PRNGKey(1))
 
-    def model(data, labels):
+    def model(data=None, labels=None):
         coefs = numpyro.sample("coefs", dist.Normal(0, 1).expand([dim]).to_event())
-        logits = numpyro.deterministic("logits", jnp.sum(coefs * data, axis=-1))
-        with numpyro.plate("N", len(data)):
-            return numpyro.sample("obs", dist.Bernoulli(logits=logits), obs=labels)
+        if data is not None:
+            logits = numpyro.deterministic("logits", jnp.sum(coefs * data, axis=-1))
+            with numpyro.plate("N", len(data)):
+                return numpyro.sample("obs", dist.Bernoulli(logits=logits), obs=labels)
 
     adam = optim.Adam(0.01)
     rng_key_init = random.PRNGKey(1)
@@ -507,12 +508,12 @@ def create_plates(batch, subsample, full_size):
 )
 def test_autoguide_deterministic(auto_class):
     def model(y=None):
-        n = y.size if y is not None else 1
+        n, len_y = (y.size, len(y)) if y is not None else (1, 1)
 
         mu = numpyro.sample("mu", dist.Normal(0, 5))
         sigma = numpyro.param("sigma", 1, constraint=constraints.positive)
 
-        with numpyro.plate("N", len(y)):
+        with numpyro.plate("N", len_y):
             y = numpyro.sample("y", dist.Normal(mu, sigma).expand((n,)), obs=y)
         numpyro.deterministic("z", (y - mu) / sigma)
 
@@ -931,3 +932,34 @@ def create_plates():
     mf_elbo = -mf_elbo.item()
 
     assert dais_elbo > mf_elbo + 0.1
+
+
+def test_autodelta_capture_deterministic_variables():
+    def model():
+        x = numpyro.sample("x", dist.Normal())
+        numpyro.deterministic("x2", x**2)
+
+    guide = AutoDelta(model)
+    svi = SVI(model, guide, optim.Adam(0.01), Trace_ELBO())
+    svi_result = svi.run(random.PRNGKey(0), num_steps=1_000)
+    guide_samples = guide.sample_posterior(
+        rng_key=random.PRNGKey(1), params=svi_result.params
+    )
+    assert "x2" in guide_samples
+
+
+@pytest.mark.parametrize("shape", [(), (1,), (2, 3)])
+@pytest.mark.parametrize("sample_shape", [(), (1,), (2, 3)])
+def test_autodelta_sample_posterior_with_sample_shape(shape, sample_shape):
+    def model():
+        x = numpyro.sample("x", dist.Normal().expand(shape))
+        numpyro.deterministic("x2", x**2)
+
+    guide = AutoDelta(model)
+    svi = SVI(model, guide, optim.Adam(0.01), Trace_ELBO())
+    svi_result = svi.run(random.PRNGKey(0), num_steps=1_000)
+    guide_samples = guide.sample_posterior(
+        rng_key=random.PRNGKey(1), params=svi_result.params, sample_shape=sample_shape
+    )
+    assert guide_samples["x"].shape == sample_shape + shape
+    assert guide_samples["x2"].shape == sample_shape + shape