Repository: arviz-devs/arviz

Issue:
`az.from_numpyro` crashes when running MCMC with `thinning != 1`
**Describe the bug**
When using `az.from_numpyro` to import an MCMC run from `numpyro`, a `ValueError` is thrown if the MCMC was run with, e.g., `thinning=2`.

**To Reproduce**
Running
```python
import arviz as az
import numpy as np
import numpyro
import numpyro.distributions as dist
from jax import random
from numpyro.infer import MCMC, NUTS


data = np.random.normal(10, 3, size=100)

def model(data):
    numpyro.sample(
        'x',
        dist.Normal(
            numpyro.sample('loc', dist.Uniform(0, 20)),
            numpyro.sample('scale', dist.Uniform(0, 20)),
        ),
        obs=data,
    )


kernel = NUTS(model)
mcmc = MCMC(NUTS(model), 100, 200, thinning=2)

mcmc.run(random.PRNGKey(0), data=data)
mcmc.print_summary()

az.from_numpyro(mcmc)  # crash
```
will crash with
```python
ValueError: cannot reshape array of size 10000 into shape (1,200,100)
```

Setting `thinning=1` will prevent the crash.

**Expected behavior**
I would expect `az.from_numpyro(mcmc)` to not crash.

**Additional context**
I am using arviz `0.11.2` and numpyro `0.6.0`.



Discussion:
Thanks for reporting, it looks like the culprit is this line:

https://github.com/arviz-devs/arviz/blob/59575274a5921acae24930733b7125c96ff8e97b/arviz/data/io_numpyro.py#L94

which gathers info from the numpyro object in order to be able to reshape the samples accordingly afterwards (which is where the error is triggered). Taking thinning into account when defining `self.ndraws` should fix the issue. It looks like `thinning` is stored in the same way as `num_chains` and `num_draws`: https://github.com/pyro-ppl/numpyro/blob/master/numpyro/infer/mcmc.py#L245-L250

Do you want to submit a PR for this @kpj ? I can help f you decide to do so, otherwise I'll label it as `beginner` and wait for someone to work on this
Thanks for the information! I'd be happy to submit a PR.

Would the fix be as easy as dividing `self.ndraws` by `posterior.thinning` in the line after https://github.com/arviz-devs/arviz/blob/59575274a5921acae24930733b7125c96ff8e97b/arviz/data/io_numpyro.py#L94?
> Would the fix be as easy as dividing self.ndraws by posterior.thinning in the line after 

I think so

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
@@ -91,7 +91,10 @@ def arbitrary_element(dct):
                     for i, v in enumerate(tree_flatten_samples)
                 }
             self._samples = samples
-            self.nchains, self.ndraws = posterior.num_chains, posterior.num_samples
+            self.nchains, self.ndraws = (
+                posterior.num_chains,
+                posterior.num_samples // posterior.thinning,
+            )
             self.model = self.posterior.sampler.model
             # model arguments and keyword arguments
             self._args = self.posterior._args  # pylint: disable=protected-access


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
@@ -227,3 +227,29 @@ def test_inference_data_num_chains(self, predictions_data, chains):
         inference_data = from_numpyro(predictions=predictions, num_chains=chains)
         nchains = inference_data.predictions.dims["chain"]
         assert nchains == chains
+
+    @pytest.mark.parametrize("nchains", [1, 2])
+    @pytest.mark.parametrize("thin", [1, 2, 3, 5, 10])
+    def test_mcmc_with_thinning(self, nchains, thin):
+        import numpyro
+        import numpyro.distributions as dist
+        from numpyro.infer import MCMC, NUTS
+
+        x = np.random.normal(10, 3, size=100)
+
+        def model(x):
+            numpyro.sample(
+                "x",
+                dist.Normal(
+                    numpyro.sample("loc", dist.Uniform(0, 20)),
+                    numpyro.sample("scale", dist.Uniform(0, 20)),
+                ),
+                obs=x,
+            )
+
+        nuts_kernel = NUTS(model)
+        mcmc = MCMC(nuts_kernel, num_warmup=100, num_samples=400, num_chains=nchains, thinning=thin)
+        mcmc.run(PRNGKey(0), x=x)
+
+        inference_data = from_numpyro(mcmc)
+        assert inference_data.posterior["loc"].shape == (nchains, 400 // thin)