Numpyro SVI conversion

[kylejcaron]

Tell us about it

It would be convenient to have a way to convert the output of a numpyro model fit with variational inference to an arviz InferenceData object - this would be analagous to pymc, where they also use InferenceData for their VI inference, and set chains=1.

The downside of this would be that many ArviZ diagnostics are meant for HMC, not VI and this could confuse inexperienced users

Thoughts on implementation

This might be tricky because while HMC stores everything needed for arviz, the SVI object does not - the outputted SVIRunResult (a namedtuple that contains the fitted params) is also needed along with the original SVI object (which contains the model and guide functions).

Possible Option 1

az.from_numpyro(
    svi,
    svi_result=svi_result,
    prior=...,
    posterior_predictive=...,
    coords=...,
  ...
)

Possible Option 2:

az.from_numpyro_svi(
   svi,
   svi_result=svi_result,
   prior=...,
   posterior_predictive=...,
   coords=...,
   ...
)

Possible Option 3:

az.from_numpyro(
    SVIWrapper(svi, result=svi_result), # mimics the mcmc structure for getting samples
    prior=...,
    posterior_predictive=...,
    coords=...,
  ...
)

[OriolAbril]

I think I'd go for option 2, probably implemented directly in arviz-base. There it is possible to customize the "sample_dims" so it should be possible to set that to sample instead of [chain, draw] and avoid the dummy chain variable. We might have to think a bit more about default behaviour if we go down this route though.

[kylejcaron]

I think I'd go for option 2, probably implemented directly in arviz-base. There it is possible to customize the "sample_dims" so it should be possible to set that to sample instead of [chain, draw] and avoid the dummy chain variable. We might have to think a bit more about default behaviour if we go down this route though.

May tackle it this week! Should I also migrate the latest NumPyroConverter (that can infer dims automatically from a model) from arviz into arviz-base as well? I'm guessing based on #2088 that eventually the methods in arviz will instead reference arviz-base?

[OriolAbril]

That would be amazing, thanks! And yes to everything, you might also find https://python.arviz.org/en/latest/user_guide/migration_guide.html useful (and I should add a link to it in the other issue too)