changed input for form_numpyro_svi from guide to SVI instance

Author: kylejcaron

external_tests/helpers.py

@@ -169,7 +169,7 @@ def numpyro_schools_model_svi(data, draws, chains):
     guide = AutoNormal(_numpyro_noncentered_model, init_loc_fn=init_to_sample())
     svi = SVI(_numpyro_noncentered_model, guide=guide, optim=Adam(0.05), loss=Trace_ELBO())
     svi_result = svi.run(PRNGKey(0), 4000, **data)
-    return {"guide": guide, "svi_result": svi_result, "model_kwargs": data}
+    return {"svi": svi, "svi_result": svi_result, "model_kwargs": data}
 
 
 def numpyro_schools_model_svi_custom_guide(data, draws, chains):
@@ -182,10 +182,9 @@ def numpyro_schools_model_svi_custom_guide(data, draws, chains):
     svi = SVI(_numpyro_noncentered_model, guide=guide, optim=Adam(0.05), loss=Trace_ELBO())
     svi_result = svi.run(PRNGKey(0), 4000, **data)
     return {
-        "guide": guide,
+        "svi": svi,
         "svi_result": svi_result,
         "model_kwargs": data,
-        "model": _numpyro_noncentered_model,
     }
 
 

external_tests/test_numpyro.py

@@ -564,7 +564,7 @@ def _run_inference(self, model, svi, guide_fn=autoguide.AutoNormal):
             svi = SVI(model, guide=guide, optim=Adam(0.05), loss=Trace_ELBO())
             svi_result = svi.run(PRNGKey(0), 10)
             return {
-                "guide": guide,
+                "svi": svi,
                 "svi_result": svi_result,
                 "model": None if is_autoguide else model,
             }

src/arviz_base/io_numpyro.py

@@ -16,16 +16,15 @@ class SVIWrapper:
 
     def __init__(
         self,
-        guide,
+        svi,
         *,
         svi_result,
         model_args=None,
         model_kwargs=None,
         num_samples: int = 1000,
-        model=None,
         thinning: int = 1,
     ):
-        self.guide = guide
+        self.svi = svi
         self.svi_result = svi_result
         self._args = model_args or tuple()
         self._kwargs = model_kwargs or dict()
@@ -34,7 +33,6 @@ def __init__(
         self.num_chains = 0
         self.sample_dims = ["samples"]
         self.kind = "svi"
-        self.model = model
 
     def get_samples(self, seed=None, **kwargs):
         """Mimics mcmc.get_samples()."""
@@ -43,8 +41,8 @@ def get_samples(self, seed=None, **kwargs):
         from numpyro.infer.autoguide import AutoGuide
 
         key = jax.random.PRNGKey(seed or 0)
-        if isinstance(self.guide, AutoGuide):
-            return self.guide.sample_posterior(
+        if isinstance(self.svi.guide, AutoGuide):
+            return self.svi.guide.sample_posterior(
                 key,
                 self.svi_result.params,
                 *self._args,
@@ -53,7 +51,7 @@ def get_samples(self, seed=None, **kwargs):
             )
         # if a custom guide is provided, sample by hand
         predictive = Predictive(
-            self.guide, params=self.svi_result.params, num_samples=self.num_samples
+            self.svi.guide, params=self.svi_result.params, num_samples=self.num_samples
         )
         samples = predictive(key, *self._args, **self._kwargs)
         return samples
@@ -70,7 +68,7 @@ def __init__(self, model):
             def model(self):
                 return self._model
 
-        return Sampler(getattr(self.guide, "model", self.model))
+        return Sampler(getattr(self.svi.guide, "model", self.svi.model))
 
     def get_extra_fields(self, **kwargs):
         """Mimics mcmc.get_extra_fields()."""
@@ -623,7 +621,7 @@ def from_numpyro(
 
 
 def from_numpyro_svi(
-    guide,
+    svi,
     svi_result,
     *,
     model_args=None,
@@ -663,8 +661,8 @@ def from_numpyro_svi(
 
     Parameters
     ----------
-    guide : numpyro.infer.autoguide.AutoGuide or callable
-        Guide function for a numpyro SVI model. Can be an autoguide or custom guide.
+    guide : numpyro.infer.svi.SVI
+        Numpyro SVI instance used for fitting the model.
     svi_result : numpyro.infer.svi.SVIRunResult
         SVI results from a fitted model.
     model_args : tuple, optional
@@ -694,20 +692,17 @@ def from_numpyro_svi(
         their coordinates.
     num_chains : int, default 1
         Number of chains used for sampling. Ignored if posterior is present.
-    model : callable, optional
-        Model function, only needed for a custom guide function
 
     Returns
     -------
     DataTree
     """
     posterior = SVIWrapper(
-        guide,
+        svi,
         svi_result=svi_result,
         model_args=model_args,
         model_kwargs=model_kwargs,
         num_samples=num_samples,
-        model=model,
     )
     with rc_context(rc={"data.sample_dims": ["samples"]}):
         return NumPyroConverter(

src/arviz_base/io_numpyro.pyi

@@ -18,13 +18,12 @@ from arviz_base.utils import expand_dims
 class SVIWrapper:
     def __init__(
         self,
-        guide,
+        svi,
         *,
         svi_result,
         model_args=...,
         model_kwargs=...,
         num_samples: int = ...,
-        model=...,
         thinning: int = ...,
     ) -> None: ...
     def get_samples(self, seed=..., **kwargs) -> None: ...
@@ -96,7 +95,7 @@ def from_numpyro(
     num_chains: int = ...,
 ) -> xarray.DataTree: ...
 def from_numpyro_svi(
-    guide: numpyro.infer.autoguide.AutoGuide | Callable,
+    svi,
     svi_result: numpyro.infer.svi.SVIRunResult,
     *,
     model_args: tuple | None = ...,
@@ -112,6 +111,6 @@ def from_numpyro_svi(
     dims: dict[str, list[str]] | None = ...,
     pred_dims: dict | None = ...,
     extra_event_dims: dict | None = ...,
-    model: Callable | None = ...,
+    model=...,
     num_samples: int = ...,
 ) -> xarray.DataTree: ...