internlm/SWE-Fixer-Train-110K huggingface dataset:
Repository: pyro-ppl/numpyro

Issue:
Bernoulli returns a Boolean instead of an integer
Currently, in Bernoulli distribution, we sample `boolean` arrays but using `arange(2)` for enumerate support. We should use 0, 1 values in `sample` method for consistency.


Discussion:

Resolving patch:
diff --git a/numpyro/distributions/discrete.py b/numpyro/distributions/discrete.py
--- a/numpyro/distributions/discrete.py
+++ b/numpyro/distributions/discrete.py
@@ -82,7 +82,8 @@ def __init__(self, probs, validate_args=None):
         super(BernoulliProbs, self).__init__(batch_shape=jnp.shape(self.probs), validate_args=validate_args)
 
     def sample(self, key, sample_shape=()):
-        return random.bernoulli(key, self.probs, shape=sample_shape + self.batch_shape)
+        samples = random.bernoulli(key, self.probs, shape=sample_shape + self.batch_shape)
+        return samples.astype(jnp.result_type(samples, int))
 
     @validate_sample
     def log_prob(self, value):
@@ -114,7 +115,8 @@ def __init__(self, logits=None, validate_args=None):
         super(BernoulliLogits, self).__init__(batch_shape=jnp.shape(self.logits), validate_args=validate_args)
 
     def sample(self, key, sample_shape=()):
-        return random.bernoulli(key, self.probs, shape=sample_shape + self.batch_shape)
+        samples = random.bernoulli(key, self.probs, shape=sample_shape + self.batch_shape)
+        return samples.astype(jnp.result_type(samples, int))
 
     @validate_sample
     def log_prob(self, value):
diff --git a/numpyro/distributions/distribution.py b/numpyro/distributions/distribution.py
--- a/numpyro/distributions/distribution.py
+++ b/numpyro/distributions/distribution.py
@@ -27,6 +27,7 @@
 
 from collections import OrderedDict
 from contextlib import contextmanager
+import functools
 import warnings
 
 import numpy as np
@@ -84,6 +85,10 @@ def __call__(cls, *args, **kwargs):
                 return result
         return super().__call__(*args, **kwargs)
 
+    @property
+    def __wrapped__(cls):
+        return functools.partial(cls.__init__, None)
+
 
 class Distribution(metaclass=DistributionMeta):
     """


Test patch:
Repository: pyro-ppl/numpyro

Issue:
MCMC postprocessing is not compatible with vectorized chain method
Currently, when we use `chain_method="vectorized"` for models with deterministic sites, there will be an error that happens at the end of sampling phase.

More info to come...


Discussion:

Resolving patch:
diff --git a/numpyro/infer/mcmc.py b/numpyro/infer/mcmc.py
--- a/numpyro/infer/mcmc.py
+++ b/numpyro/infer/mcmc.py
@@ -7,7 +7,7 @@
 import os
 import warnings
 
-from jax import device_put, jit, lax, pmap, random
+from jax import device_put, jit, lax, pmap, random, vmap
 from jax.core import Tracer
 from jax.interpreters.xla import DeviceArray
 from jax.lib import xla_bridge
@@ -342,6 +342,8 @@ def _single_chain_mcmc(self, init, args, kwargs, collect_fields):
         # so we only need to filter out the case site_value.shape[0] == 0
         # (which happens when lower_idx==upper_idx)
         if len(site_values) > 0 and jnp.shape(site_values[0])[0] > 0:
+            if self.chain_method == "vectorized" and self.num_chains > 1:
+                postprocess_fn = vmap(postprocess_fn)
             states[self._sample_field] = lax.map(postprocess_fn, states[self._sample_field])
         return states, last_state
 


Test patch:
diff --git a/test/test_mcmc.py b/test/test_mcmc.py
--- a/test/test_mcmc.py
+++ b/test/test_mcmc.py
@@ -366,6 +366,7 @@ def test_chain(use_init_params, chain_method):
     def model(labels):
         coefs = numpyro.sample('coefs', dist.Normal(jnp.zeros(dim), jnp.ones(dim)))
         logits = jnp.sum(coefs * data, axis=-1)
+        numpyro.deterministic("logits", logits)
         return numpyro.sample('obs', dist.Bernoulli(logits=logits), obs=labels)
 
     kernel = NUTS(model=model)
Repository: pyro-ppl/numpyro

Issue:
Update docstring of Neal's funnel example
We have updated [funnel](https://github.com/pyro-ppl/numpyro/blob/master/examples/funnel.py) example to use `reparam` handler, but the docstring is not updated yet.


Discussion:

Resolving patch:
diff --git a/examples/funnel.py b/examples/funnel.py
--- a/examples/funnel.py
+++ b/examples/funnel.py
@@ -6,18 +6,15 @@
 ======================
 
 This example, which is adapted from [1], illustrates how to leverage non-centered
-parameterization using the class :class:`numpyro.distributions.TransformedDistribution`.
+parameterization using the :class:`~numpyro.handlers.reparam` handler.
 We will examine the difference between two types of parameterizations on the
 10-dimensional Neal's funnel distribution. As we will see, HMC gets trouble at
 the neck of the funnel if centered parameterization is used. On the contrary,
 the problem can be solved by using non-centered parameterization.
 
-Using non-centered parameterization through TransformedDistribution in NumPyro
-has the same effect as the automatic reparameterisation technique introduced in
-[2]. However, in [2], users need to implement a (non-trivial) reparameterization
-rule for each type of transform. Instead, in NumPyro the only requirement to let
-inference algorithms know to do reparameterization automatically is to declare
-the random variable as a transformed distribution.
+Using non-centered parameterization through :class:`~numpyro.infer.reparam.LocScaleReparam`
+or :class:`~numpyro.infer.reparam.TransformReparam` in NumPyro has the same effect as
+the automatic reparameterisation technique introduced in [2].
 
 **References:**
 
@@ -36,6 +33,7 @@
 
 import numpyro
 import numpyro.distributions as dist
+from numpyro.handlers import reparam
 from numpyro.infer import MCMC, NUTS, Predictive
 from numpyro.infer.reparam import LocScaleReparam
 
@@ -45,10 +43,7 @@ def model(dim=10):
     numpyro.sample('x', dist.Normal(jnp.zeros(dim - 1), jnp.exp(y / 2)))
 
 
-def reparam_model(dim=10):
-    y = numpyro.sample('y', dist.Normal(0, 3))
-    with numpyro.handlers.reparam(config={'x': LocScaleReparam(0)}):
-        numpyro.sample('x', dist.Normal(jnp.zeros(dim - 1), jnp.exp(y / 2)))
+reparam_model = reparam(model, config={'x': LocScaleReparam(0)})
 
 
 def run_inference(model, args, rng_key):
@@ -56,7 +51,7 @@ def run_inference(model, args, rng_key):
     mcmc = MCMC(kernel, args.num_warmup, args.num_samples, num_chains=args.num_chains,
                 progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True)
     mcmc.run(rng_key)
-    mcmc.print_summary()
+    mcmc.print_summary(exclude_deterministic=False)
     return mcmc.get_samples()
 
 


Test patch:
Repository: pyro-ppl/numpyro

Issue:
improve gradient computation of eucledian kinetic energy
I have been wanted to do this for a while. Given the recent improvement in JAX regarding custom derivatives, it is time to make this enhancement. Currently, `kinetic_energy` is defined as `x.T @ M @ x`, which has an analytical grad formula (we will use this instead of relying on `grad`). This is just a small optimization but would be useful for the `dense_mass` case.


Discussion:

Resolving patch:
diff --git a/numpyro/infer/hmc_util.py b/numpyro/infer/hmc_util.py
--- a/numpyro/infer/hmc_util.py
+++ b/numpyro/infer/hmc_util.py
@@ -170,6 +170,13 @@ def final_fn(state, regularize=False):
     return init_fn, update_fn, final_fn
 
 
+def _kinetic_grad(kinetic_fn, inverse_mass_matrix, r):
+    if hasattr(kinetic_fn, "_kinetic_grad"):
+        return kinetic_fn._kinetic_grad(inverse_mass_matrix, r)
+    else:
+        return grad(kinetic_fn, argnums=1)(inverse_mass_matrix, r)
+
+
 def velocity_verlet(potential_fn, kinetic_fn):
     r"""
     Second order symplectic integrator that uses the velocity verlet algorithm
@@ -204,7 +211,7 @@ def update_fn(step_size, inverse_mass_matrix, state):
         """
         z, r, _, z_grad = state
         r = tree_multimap(lambda r, z_grad: r - 0.5 * step_size * z_grad, r, z_grad)  # r(n+1/2)
-        r_grad = grad(kinetic_fn, argnums=1)(inverse_mass_matrix, r)
+        r_grad = _kinetic_grad(kinetic_fn, inverse_mass_matrix, r)
         z = tree_multimap(lambda z, r_grad: z + step_size * r_grad, z, r_grad)  # z(n+1)
         potential_energy, z_grad = value_and_grad(potential_fn)(z)
         r = tree_multimap(lambda r, z_grad: r - 0.5 * step_size * z_grad, r, z_grad)  # r(n+1)
@@ -727,6 +734,20 @@ def euclidean_kinetic_energy(inverse_mass_matrix, r):
     return 0.5 * jnp.dot(v, r)
 
 
+def _euclidean_kinetic_energy_grad(inverse_mass_matrix, r):
+    r, unravel_fn = ravel_pytree(r)
+
+    if inverse_mass_matrix.ndim == 2:
+        v = jnp.matmul(inverse_mass_matrix, r)
+    elif inverse_mass_matrix.ndim == 1:
+        v = jnp.multiply(inverse_mass_matrix, r)
+
+    return unravel_fn(v)
+
+
+euclidean_kinetic_energy._kinetic_grad = _euclidean_kinetic_energy_grad
+
+
 def consensus(subposteriors, num_draws=None, diagonal=False, rng_key=None):
     """
     Merges subposteriors following consensus Monte Carlo algorithm.


Test patch:
Repository: pyro-ppl/numpyro

Issue:
ZeroDivisionError when running SVI with num_steps < 20
Reproducible code
```python
import jax
import numpyro

def model():
    pass

def guide():
    pass

svi = numpyro.infer.SVI(model, guide, numpyro.optim.Adam(1), numpyro.infer.Trace_ELBO())
svi.run(jax.random.PRNGKey(0), 10)
```

Proposed fix: setting `batch = max(num_steps // 20, 1)` in [SVI.run](https://github.com/pyro-ppl/numpyro/blob/aee945895dff9ffb6a12637a9bfbddb1d8c20363/numpyro/infer/svi.py#L179) to guarantee that `batch` is positive


Discussion:

Resolving patch:
diff --git a/numpyro/infer/svi.py b/numpyro/infer/svi.py
--- a/numpyro/infer/svi.py
+++ b/numpyro/infer/svi.py
@@ -176,15 +176,15 @@ def body_fn(svi_state, carry):
         if progress_bar:
             losses = []
             with tqdm.trange(1, num_steps + 1) as t:
-                batch = num_steps // 20
+                batch = max(num_steps // 20, 1)
                 for i in t:
                     svi_state, loss = jit(body_fn)(svi_state, None)
+                    losses.append(loss)
                     if i % batch == 0:
                         avg_loss = sum(losses[i-batch:]) / batch
                         t.set_postfix_str("init loss: {:.4f}, avg. loss [{}-{}]: {:.4f}"
                                           .format(losses[0], i - batch + 1, i, avg_loss),
                                           refresh=False)
-                    losses.append(loss)
             losses = jnp.stack(losses)
         else:
             svi_state, losses = lax.scan(body_fn, svi_state, None, length=num_steps)


Test patch:
diff --git a/test/test_svi.py b/test/test_svi.py
--- a/test/test_svi.py
+++ b/test/test_svi.py
@@ -174,3 +174,15 @@ def guide():
     assert jnp.isfinite(actual_loss)
     expected_loss = x_guide.log_prob(x) - x_prior.log_prob(x)
     assert_allclose(actual_loss, expected_loss)
+
+
+@pytest.mark.parametrize("num_steps", [10, 30, 50])
+def test_run_with_small_num_steps(num_steps):
+    def model():
+        pass
+
+    def guide():
+        pass
+
+    svi = SVI(model, guide, optim.Adam(1), Trace_ELBO())
+    svi.run(random.PRNGKey(0), num_steps)
Repository: pyro-ppl/numpyro

Issue:
test_compile fails locally but the error is silent in CI
For some reason, the model is traced more than needed. This will a bit affect large models.

```python
FAILED test/test_compile.py::test_mcmc_one_chain[True-True] - assert 8 == 6
FAILED test/test_compile.py::test_mcmc_one_chain[True-False] - assert 7 == 5
FAILED test/test_compile.py::test_mcmc_one_chain[False-True] - assert 5 == 4
FAILED test/test_compile.py::test_mcmc_one_chain[False-False] - assert 4 == 3
```

I have no idea why some errors (like in doctest or this one) are silent on CI...


Discussion:

Resolving patch:
diff --git a/numpyro/infer/barker.py b/numpyro/infer/barker.py
--- a/numpyro/infer/barker.py
+++ b/numpyro/infer/barker.py
@@ -165,7 +165,9 @@ def init(self, rng_key, num_warmup, init_params, model_args, model_kwargs):
         size = len(ravel_pytree(init_params)[0])
         wa_state = wa_init(None, rng_key_wa, self._step_size, mass_matrix_size=size)
         wa_state = wa_state._replace(rng_key=None)
-        return jax.device_put(BarkerMHState(0, init_params, pe, grad, 0., 0., wa_state, rng_key))
+        init_state = BarkerMHState(jnp.array(0), init_params, pe, grad, jnp.array(0.),
+                                   jnp.array(0.), wa_state, rng_key)
+        return jax.device_put(init_state)
 
     def postprocess_fn(self, args, kwargs):
         if self._postprocess_fn is None:
diff --git a/numpyro/infer/hmc.py b/numpyro/infer/hmc.py
--- a/numpyro/infer/hmc.py
+++ b/numpyro/infer/hmc.py
@@ -245,8 +245,8 @@ def init_kernel(init_params,
         vv_init, vv_update = velocity_verlet(pe_fn, kinetic_fn, forward_mode_ad)
         vv_state = vv_init(z, r, potential_energy=pe, z_grad=z_grad)
         energy = kinetic_fn(wa_state.inverse_mass_matrix, vv_state.r)
-        hmc_state = HMCState(0, vv_state.z, vv_state.z_grad, vv_state.potential_energy, energy,
-                             0, 0., 0., False, wa_state, rng_key_hmc)
+        hmc_state = HMCState(jnp.array(0), vv_state.z, vv_state.z_grad, vv_state.potential_energy, energy,
+                             jnp.array(0), jnp.array(0.), jnp.array(0.), jnp.array(False), wa_state, rng_key_hmc)
         return device_put(hmc_state)
 
     def _hmc_next(step_size, inverse_mass_matrix, vv_state,
diff --git a/numpyro/infer/sa.py b/numpyro/infer/sa.py
--- a/numpyro/infer/sa.py
+++ b/numpyro/infer/sa.py
@@ -135,7 +135,8 @@ def init_kernel(init_params,
         k = random.categorical(rng_key_z, jnp.zeros(zs.shape[0]))
         z = unravel_fn(zs[k])
         pe = pes[k]
-        sa_state = SAState(0, z, pe, 0., 0., False, adapt_state, rng_key_sa)
+        sa_state = SAState(jnp.array(0), z, pe, jnp.array(0.), jnp.array(0.), jnp.array(False),
+                           adapt_state, rng_key_sa)
         return device_put(sa_state)
 
     def sample_kernel(sa_state, model_args=(), model_kwargs=None):


Test patch:
Repository: pyro-ppl/numpyro

Issue:
Add required packages at the top of tutorials to run on colab



Discussion:

Resolving patch:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -55,9 +55,12 @@
         "dev": [
             "dm-haiku",
             "flax",
+            # TODO: bump funsor version before the release
             "funsor @ git+https://github.com/pyro-ppl/funsor.git@d5574988665dd822ec64e41f2b54b9dc929959dc",
             "graphviz",
-            "tensorflow_probability",
+            # TODO: change this to tensorflow_probability>0.12.1 when the next version
+            # of tfp is released. The current release is not compatible with jax>=0.2.12.
+            "tfp-nightly",
         ],
         "examples": ["arviz", "jupyter", "matplotlib", "pandas", "seaborn"],
     },


Test patch:
Repository: pyro-ppl/numpyro

Issue:
dist.DirichletMultinomial can't handle zero_total count
```
import jax
import jax.numpy as jnp
from numpyro.distributions import DirichletMultinomial

c = jnp.ones(3)
t = jnp.array(0)  # 1 is ok
d = DirichletMultinomial(concentration=c, total_count=t)
z = d.sample(key=jax.random.PRNGKey(0))
assert z.shape == (3,)
```

expected behavior: `total_count=0` is ok


Discussion:

Resolving patch:
diff --git a/numpyro/distributions/util.py b/numpyro/distributions/util.py
--- a/numpyro/distributions/util.py
+++ b/numpyro/distributions/util.py
@@ -215,6 +215,8 @@ def _multinomial(key, p, n, n_max, shape=()):
         n = jnp.broadcast_to(n, broadcast_shape)
         p = jnp.broadcast_to(p, broadcast_shape + jnp.shape(p)[-1:])
     shape = shape or p.shape[:-1]
+    if n_max == 0:
+        return jnp.zeros(shape + p.shape[-1:], dtype=jnp.result_type(int))
     # get indices from categorical distribution then gather the result
     indices = categorical(key, p, (n_max,) + shape)
     # mask out values when counts is heterogeneous
@@ -244,7 +246,10 @@ def _multinomial(key, p, n, n_max, shape=()):
 
 
 def multinomial(key, p, n, shape=()):
-    n_max = int(jnp.max(n))
+    assert not isinstance(
+        n, jax.core.Tracer
+    ), "The total count parameter `n` should not be a jax abstract array."
+    n_max = int(np.max(jax.device_get(n)))
     return _multinomial(key, p, n, n_max, shape)
 
 


Test patch:
diff --git a/test/test_distributions_util.py b/test/test_distributions_util.py
--- a/test/test_distributions_util.py
+++ b/test/test_distributions_util.py
@@ -91,6 +91,18 @@ def test_multinomial_shape(p, shape):
     assert jnp.shape(multinomial(rng_key, p, n, shape)) == expected_shape
 
 
+@pytest.mark.parametrize("n", [0, 1, np.array([0, 0]), np.array([2, 1, 0])])
+@pytest.mark.parametrize("device_array", [True, False])
+def test_multinomial_inhomogeneous(n, device_array):
+    if device_array:
+        n = jnp.asarray(n)
+
+    p = jnp.array([0.5, 0.5])
+    x = multinomial(random.PRNGKey(0), p, n)
+    assert x.shape == jnp.shape(n) + jnp.shape(p)
+    assert_allclose(x.sum(-1), n)
+
+
 @pytest.mark.parametrize("p", [np.array([0.2, 0.3, 0.5]), np.array([0.8, 0.1, 0.1])])
 @pytest.mark.parametrize("n", [10000, np.array([10000, 20000])])
 def test_multinomial_stats(p, n):
Repository: pyro-ppl/numpyro

Issue:
New versions of sphinx and jinja2 break docs linting
This is observed by @tcbegley in #1034


Discussion:
The sphinx issue has been fixed in [nbsphinx==0.8.5](https://github.com/spatialaudio/nbsphinx/pull/562). Also, in `nbsphinx==0.8.5`, jinja2 version is [forced to be smaller than 3](https://github.com/spatialaudio/nbsphinx/issues/563) so as long as we use the latest nbsphinx version, it will be fine.

Resolving patch:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -39,9 +39,8 @@
     extras_require={
         "doc": [
             "ipython",  # sphinx needs this to render codes
-            "jinja2<3.0.0",
-            "nbsphinx",
-            "sphinx<4.0.0",
+            "nbsphinx>=0.8.5",
+            "sphinx",
             "sphinx_rtd_theme",
             "sphinx-gallery",
         ],


Test patch:
Repository: pyro-ppl/numpyro

Issue:
[Docs] Syntax highlighting broken for notebook examples / tutorials
Syntax highlighting doesn't seem to work on any notebook based examples tutorials in the deployed documentation. See e.g. http://num.pyro.ai/en/stable/tutorials/bayesian_regression.html

Interestingly I don't have the same problems if I build the documentation locally (following steps in a fresh virtual environment)

```
pip install .[doc]
pip install -r docs/requirements.txt
make docs
python -m http.server --directory docs/build/html
```

I can't see from the repo how the docs are built / deployed, so I suspect there's something in that process that disagrees with the above steps? In particular I think the `ipython` dependency is key, is that maybe not getting picked up in the build of the deployed docs?


Discussion:
In the [latest build](https://readthedocs.org/projects/numpyro/builds/14369063/), readthedocs uses
```python
pip install --upgrade --no-cache-dir mock==1.0.1 pillow==5.4.1 alabaster>=0.7,<0.8,!=0.7.5 commonmark==0.8.1 recommonmark==0.5.0 sphinx<2 sphinx-rtd-theme<0.5 readt
```
Looking at [rtd documentation](https://docs.readthedocs.io/en/stable/guides/reproducible-builds.html#pinning-dependencies) (the docs point to this [search](https://github.com/readthedocs/readthedocs-sphinx-search) package, which is pretty cool and convenient - we should use it), I think that we need to pin
```
nbsphinx==0.8.6
sphinx==4.1.2
sphinx_rtd_theme==0.5.2
# readthedocs-sphinx-search==0.1.0
```
in `docs/requirements.txt` file. But IIRC, the latest versions do not give the correct `Open in Colab`/`Open in Github` links... I can't remember exactly. We can try those latest versions again. Do you want to make a PR?

Resolving patch:
diff --git a/docs/source/conf.py b/docs/source/conf.py
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -73,6 +73,7 @@
     "sphinx.ext.mathjax",
     "sphinx.ext.viewcode",
     "sphinx_gallery.gen_gallery",
+    "sphinx_search.extension",
 ]
 
 # Enable documentation inheritance
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -47,6 +47,7 @@
         "doc": [
             "ipython",  # sphinx needs this to render codes
             "nbsphinx>=0.8.5",
+            "readthedocs-sphinx-search==0.1.0",
             "sphinx",
             "sphinx_rtd_theme",
             "sphinx-gallery",


Test patch:
Repository: pyro-ppl/numpyro

Issue:
Use Gauss-Newton hessian matrix in laplace approximation
Currently, in AutoLaplaceApproximation, we use `jax.hessian` to compute Hessian matrix. But the hessian matrix is not always positive definite. So it would be nice to use Gaussian-Newton hessian matrix when [scale_tril = Nan](https://github.com/pyro-ppl/numpyro/blob/8bb94f170de3f6c276fe61e4c92cd4e21de70a4b/numpyro/infer/autoguide.py#L898).


Discussion:

Resolving patch:
diff --git a/numpyro/infer/autoguide.py b/numpyro/infer/autoguide.py
--- a/numpyro/infer/autoguide.py
+++ b/numpyro/infer/autoguide.py
@@ -889,7 +889,7 @@ def loss_fn(z):
                 warnings.warn(
                     "Hessian of log posterior at the MAP point is singular. Posterior"
                     " samples from AutoLaplaceApproxmiation will be constant (equal to"
-                    " the MAP point)."
+                    " the MAP point). Please consider using an AutoNormal guide."
                 )
         scale_tril = jnp.where(jnp.isnan(scale_tril), 0.0, scale_tril)
         return LowerCholeskyAffine(loc, scale_tril)


Test patch:
Repository: pyro-ppl/numpyro

Issue:
SA kernel missing model attribute
Trying to load an MCMC run using an SA kernel into `arviz`, I ran into this issue: 

```
def test_model(...)
    ...
kernel = SA(test_model)
mcmc_test = MCMC(kernel, ...)
mcmc_test.run(...)

data_test = az.from_numpyro(mcmc_test)
```

which raises...
```---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
/tmp/ipykernel_20194/2118154136.py in <module>
----> 1 data_test = az.from_numpyro(mcmc_test)

~/miniconda3/envs/refit_fvs/lib/python3.9/site-packages/arviz/data/io_numpyro.py in from_numpyro(posterior, prior, posterior_predictive, predictions, constant_data, predictions_constant_data, coords, dims, pred_dims, num_chains)
    331         Number of chains used for sampling. Ignored if posterior is present.
    332     """
--> 333     return NumPyroConverter(
    334         posterior=posterior,
    335         prior=prior,

~/miniconda3/envs/refit_fvs/lib/python3.9/site-packages/arviz/data/io_numpyro.py in __init__(self, posterior, prior, posterior_predictive, predictions, constant_data, predictions_constant_data, coords, dims, pred_dims, num_chains)
     91             self._samples = samples
     92             self.nchains, self.ndraws = posterior.num_chains, posterior.num_samples
---> 93             self.model = self.posterior.sampler.model
     94             # model arguments and keyword arguments
     95             self._args = self.posterior._args  # pylint: disable=protected-access

AttributeError: 'SA' object has no attribute 'model'
```

Looking at the source code and noticing that this works for NUTS and HMC, the missing piece in the SA class seems to be the `model` property... going to prepare a PR to add it in the same manner as it's found in NUTS and HMC classes. 


Discussion:

Resolving patch:
diff --git a/numpyro/infer/sa.py b/numpyro/infer/sa.py
--- a/numpyro/infer/sa.py
+++ b/numpyro/infer/sa.py
@@ -365,6 +365,10 @@ def init(
             self._sample_fn = sample_fn
         return init_state
 
+    @property
+    def model(self):
+        return self._model
+
     @property
     def sample_field(self):
         return "z"


Test patch: