Fix az.from_numpyro crash when running MCMC with thinning > 1

[kpj]

Description

This PR fixes the crash observed when running the following code:

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

Fixes #1618 (line edited by @OriolAbril to use github issue closing keywords)

Checklist

• Follows official PR format
• Includes new or updated tests to cover the new feature
• Code style correct (follows pylint and black guidelines)
• Changes are listed in changelog

[ahartikainen]

Hi, added couple of comments. Looks good.

[OriolAbril]

Thanks for the superfast PR already with tests! ❤️

[kpj]

Thanks for your great suggestions :-)

I implemented them and the CI/CD errors do not seem to come from my code.

[ahartikainen]

Ok, these are not related to your code

arviz/data/inference_data.py:1630 in public function `concat`:
        D418: Function/ Method decorated with @overload shouldn't contain a docstring
arviz/data/inference_data.py:1642 in public function `concat`:
        D418: Function/ Method decorated with @overload shouldn't contain a docstring
arviz/data/inference_data.py:1655 in public function `concat`:
        D418: Function/ Method decorated with @overload shouldn't contain a docstring
arviz/data/inference_data.py:1668 in public function `concat`:
        D418: Function/ Method decorated with @overload shouldn't contain a docstring
arviz/data/inference_data.py:1681 in public function `concat`:
        D418: Function/ Method decorated with @overload shouldn't contain a docstring

I will quickly fix in in another PR.

[ahartikainen]

Thanks for the PR