Repository: pyro-ppl/numpyro

Issue:
Improving error message when using NeutraReparam
Here is a snippet

```python
import numpy as np
import jax

import jax.numpy as jnp
import numpyro
import numpyro.distributions as dist
import numpyro.infer.autoguide as autoguide
from numpyro.infer import Predictive, SVI, Trace_ELBO,  TraceMeanField_ELBO
from numpyro.optim import Adam
from numpyro.infer.reparam import NeuTraReparam
from numpyro.handlers import seed, trace, condition

#Generate mock data
param_true = np.array([1.0, 0.0, 0.2, 0.5, 1.5])
sample_size = 1_000
sigma_e = param_true[4]          # true value of parameter error sigma
random_num_generator = np.random.RandomState(0)
xi = 5*random_num_generator.rand(sample_size)-2.5
e = random_num_generator.normal(0, sigma_e, sample_size)
yi = param_true[0] + param_true[1] * xi + param_true[2] * xi**2 + param_true[3] *xi**3# +  e  

# Simple model
def my_model(Xspls,Yspls=None):
    a0 = numpyro.sample('a0', dist.Normal(0.,10.))
    a1 = numpyro.sample('a1', dist.Normal(0.,10.))
    a2 = numpyro.sample('a2', dist.Normal(0.,10.))
    a3 = numpyro.sample('a3', dist.Normal(0.,10.))

    mu = a0 + a1*Xspls + a2*Xspls**2 + a3*Xspls**3

    return numpyro.sample('obs', dist.Normal(mu, sigma_e), obs=Yspls)

#Simple SVI
guide = autoguide.AutoMultivariateNormal(my_model, init_loc_fn=numpyro.infer.init_to_sample())
optimizer = numpyro.optim.Adam(step_size=5e-3)
svi = SVI(my_model, guide,optimizer,loss=Trace_ELBO())
svi_result = svi.run(jax.random.PRNGKey(0), 10000, Xspls=xi, Yspls=yi)

#Neutra-reparametrisation
neutra = NeuTraReparam(guide, svi_result.params)
neutra_model = neutra.reparam(my_model)

#Tracing the original model (It is working)
tr = numpyro.handlers.trace(seed(my_model, jax.random.PRNGKey(0)))
tr.get_trace(Xspls=xi)

#Tracing the neutra_model (ERROR)
tr = numpyro.handlers.trace(seed(neutra_model, jax.random.PRNGKey(0)))
tr.get_trace(Xspls=xi)
```

this will generate an error until you provide `Yspls=yi` also in the `get_trace` but the error message is not so explicit and if one can improve it it would be great. 


Discussion:

Resolving patch:
diff --git a/numpyro/infer/reparam.py b/numpyro/infer/reparam.py
--- a/numpyro/infer/reparam.py
+++ b/numpyro/infer/reparam.py
@@ -235,10 +235,12 @@ def __init__(self, guide, params):
         self._x_unconstrained = {}
 
     def _reparam_config(self, site):
-        if site["name"] in self.guide.prototype_trace and not site.get(
-            "is_observed", False
-        ):
-            return self
+        if site["name"] in self.guide.prototype_trace:
+            # We only reparam if this is an unobserved site in the guide
+            # prototype trace.
+            guide_site = self.guide.prototype_trace[site["name"]]
+            if not guide_site.get("is_observed", False):
+                return self
 
     def reparam(self, fn=None):
         return numpyro.handlers.reparam(fn, config=self._reparam_config)


Test patch:
diff --git a/test/infer/test_reparam.py b/test/infer/test_reparam.py
--- a/test/infer/test_reparam.py
+++ b/test/infer/test_reparam.py
@@ -156,7 +156,7 @@ def neals_funnel(dim):
 def dirichlet_categorical(data):
     concentration = jnp.array([1.0, 1.0, 1.0])
     p_latent = numpyro.sample("p", dist.Dirichlet(concentration))
-    with numpyro.plate("N", data.shape[0]):
+    with numpyro.plate("N", data.shape[0] if data is not None else 10):
         numpyro.sample("obs", dist.Categorical(p_latent), obs=data)
     return p_latent
 
@@ -190,7 +190,7 @@ def body_fn(i, val):
     "model, kwargs",
     [
         (neals_funnel, {"dim": 10}),
-        (dirichlet_categorical, {"data": jnp.ones(10, dtype=jnp.int32)}),
+        (dirichlet_categorical, {"data": np.ones(10, dtype=np.int32)}),
     ],
 )
 def test_reparam_log_joint(model, kwargs):
@@ -212,6 +212,19 @@ def test_reparam_log_joint(model, kwargs):
     assert_allclose(pe_transformed, pe - log_det_jacobian, rtol=2e-7)
 
 
+def test_neutra_reparam_unobserved_model():
+    model = dirichlet_categorical
+    data = jnp.ones(10, dtype=jnp.int32)
+    guide = AutoIAFNormal(model)
+    svi = SVI(model, guide, Adam(1e-3), Trace_ELBO())
+    svi_state = svi.init(random.PRNGKey(0), data)
+    params = svi.get_params(svi_state)
+    neutra = NeuTraReparam(guide, params)
+    reparam_model = neutra.reparam(model)
+    with handlers.seed(rng_seed=0):
+        reparam_model(data=None)
+
+
 @pytest.mark.parametrize("shape", [(), (4,), (3, 2)], ids=str)
 @pytest.mark.parametrize("centered", [0.0, 0.6, 1.0, None])
 @pytest.mark.parametrize("dist_type", ["Normal", "StudentT"])