Repository: arviz-devs/arviz

Issue:
NumPyroConverter makes too many assumptions about the model
The following code is not always runnable:

https://github.com/arviz-devs/arviz/blob/7ebedd2d1494f07513bc7b9ea21962bc11399a94/arviz/data/io_numpyro.py#L126-L133

This will fail for a model involving `ImproperUniform`, as sampling for it is not defined.

Would it be possible to add a keyword argument that specifies whether this automatic extraction of observations should be attempt or not, e.g. `extract_observations` or something? At the moment it's stopping me from being able to use arviz with my numpyro-model (unless I just comment out the lines above) :confused: 


Discussion:
I think adding an extra argument to skip this sounds fine, as would probably be adding a try except or even both. cc @fehiepsi 
Because this just requires a valid trace, I believe we can just use some sort of init_to_uniform strategy. I'll push a PR shortly.
> Because this just requires a valid trace, I believe we can just use some sort of init_to_uniform strategy. I'll push a PR shortly.

True, but even so, as a user, I'd prefer not always inferring this, i.e. at least allow me to turn it off. E.g. the observations might be not be present because I'm "observing" using a factor, in which case the automatically inferred observations can be misleading.
That's a good point! I think you can set `observations = None` or delete some of its keys if you want. It will give you some flexibility in extracting `log_likelihood` or something like that.
How about adding a kwarg `observations=None` and then do `if observations is None and self.model is not None` or something along those lines? This way you can choose not to extract the observations by setting `observations={}` or any empty collection + it allows you to override the inferred observations, e.g. if you're using a factor approach as I mentioned above.

Resolving patch:
diff --git a/arviz/data/io_numpyro.py b/arviz/data/io_numpyro.py
--- a/arviz/data/io_numpyro.py
+++ b/arviz/data/io_numpyro.py
@@ -124,7 +124,11 @@ def arbitrary_element(dct):
 
         observations = {}
         if self.model is not None:
-            seeded_model = numpyro.handlers.seed(self.model, jax.random.PRNGKey(0))
+            # we need to use an init strategy to generate random samples for ImproperUniform sites
+            seeded_model = numpyro.handlers.substitute(
+                numpyro.handlers.seed(self.model, jax.random.PRNGKey(0)),
+                substitute_fn=numpyro.infer.init_to_sample,
+            )
             trace = numpyro.handlers.trace(seeded_model).get_trace(*self._args, **self._kwargs)
             observations = {
                 name: site["value"]


Test patch:
diff --git a/arviz/tests/external_tests/test_data_numpyro.py b/arviz/tests/external_tests/test_data_numpyro.py
--- a/arviz/tests/external_tests/test_data_numpyro.py
+++ b/arviz/tests/external_tests/test_data_numpyro.py
@@ -253,3 +253,17 @@ def model(x):
 
         inference_data = from_numpyro(mcmc)
         assert inference_data.posterior["loc"].shape == (nchains, 400 // thin)
+
+    def test_mcmc_improper_uniform(self):
+        import numpyro
+        import numpyro.distributions as dist
+        from numpyro.infer import MCMC, NUTS
+
+        def model():
+            x = numpyro.sample("x", dist.ImproperUniform(dist.constraints.positive, (), ()))
+            return numpyro.sample("y", dist.Normal(x, 1), obs=1.0)
+
+        mcmc = MCMC(NUTS(model), num_warmup=10, num_samples=10)
+        mcmc.run(PRNGKey(0))
+        inference_data = from_numpyro(mcmc)
+        assert inference_data.observed_data