updated sample_dims

Author: kylejcaron

external_tests/test_numpyro.py

@@ -334,7 +334,7 @@ def guide():
 
         result = self._run_inference(model, svi=svi, guide_fn=guide_fn)
         from_numpyro_func = from_numpyro_svi if svi else from_numpyro
-        sample_dims = ("samples",) if svi else ("chain", "draw")
+        sample_dims = ("sample",) if svi else ("chain", "draw")
 
         inference_data = from_numpyro_func(
             **result, coords={"group1": np.arange(10), "group2": np.arange(5)}
@@ -380,7 +380,7 @@ def guide():
 
         result = self._run_inference(model, svi=svi, guide_fn=guide_fn)
         from_numpyro_func = from_numpyro_svi if svi else from_numpyro
-        sample_dims = ("samples",) if svi else ("chain", "draw")
+        sample_dims = ("sample",) if svi else ("chain", "draw")
 
         inference_data = from_numpyro_func(
             **result, coords={"group1": np.arange(10), "group2": np.arange(5)}
@@ -417,7 +417,7 @@ def guide():
 
         result = self._run_inference(model, svi=svi, guide_fn=guide_fn)
         from_numpyro_func = from_numpyro_svi if svi else from_numpyro
-        sample_dims = ("samples",) if svi else ("chain", "draw")
+        sample_dims = ("sample",) if svi else ("chain", "draw")
 
         inference_data = from_numpyro_func(**result)
         assert inference_data.posterior.param.dims == sample_dims + ("group",)
@@ -453,7 +453,7 @@ def guide():
 
         result = self._run_inference(model, svi=svi, guide_fn=guide_fn)
         from_numpyro_func = from_numpyro_svi if svi else from_numpyro
-        sample_dims = ("samples",) if svi else ("chain", "draw")
+        sample_dims = ("sample",) if svi else ("chain", "draw")
 
         inference_data = from_numpyro_func(**result, coords={"groups": np.arange(10)})
         assert inference_data.posterior.gamma.dims == sample_dims + ("groups",)
@@ -535,7 +535,7 @@ def guide():
 
         result = self._run_inference(model, svi=svi, guide_fn=guide_fn)
         from_numpyro_func = from_numpyro_svi if svi else from_numpyro
-        sample_dims = ("samples",) if svi else ("chain", "draw")
+        sample_dims = ("sample",) if svi else ("chain", "draw")
         inference_data = from_numpyro_func(
             **result, coords={"groups": np.arange(10)}, extra_event_dims={"gamma_plus1": ["groups"]}
         )
@@ -548,7 +548,7 @@ def test_predictions_infer_dims(
         inference_data = self.get_inference_data(
             data, eight_schools_params, predictions_data, predictions_params, infer_dims=True
         )
-        sample_dims = ("samples",) if isinstance(data.obj, dict) else ("chain", "draw")
+        sample_dims = ("sample",) if isinstance(data.obj, dict) else ("chain", "draw")
         assert inference_data.predictions.obs.dims == (sample_dims + ("J",))
         assert "J" in inference_data.predictions.obs.coords
 

src/arviz_base/io_numpyro.py

@@ -33,7 +33,7 @@ def __init__(
         self.num_samples = num_samples
         self.thinning = 1
         self.num_chains = 0
-        self.sample_dims = ["samples"]
+        self.sample_dims = ["sample"]
         self.kind = "svi"
 
         self.numpyro = numpyro
@@ -703,7 +703,7 @@ def from_numpyro_svi(
         model_kwargs=model_kwargs,
         num_samples=num_samples,
     )
-    with rc_context(rc={"data.sample_dims": ["samples"]}):
+    with rc_context(rc={"data.sample_dims": ["sample"]}):
         return NumPyroConverter(
             posterior=posterior,
             prior=prior,