Test observations and constant data are included in InferenceData from JAX sampling

Author: michaelosthege

pymc/backends/arviz.py

@@ -43,11 +43,8 @@
 Var = Any  # pylint: disable=invalid-name
 
 
-def find_observations(model: Optional["Model"]) -> Dict[str, Var]:
+def find_observations(model: "Model") -> Dict[str, Var]:
     """If there are observations available, return them as a dictionary."""
-    if model is None:
-        return {}
-
     observations = {}
     for obs in model.observed_RVs:
         aux_obs = getattr(obs.tag, "observations", None)
@@ -63,12 +60,9 @@ def find_observations(model: Optional["Model"]) -> Dict[str, Var]:
     return observations
 
 
-def find_constants(model: Optional["Model"]) -> Dict[str, Var]:
+def find_constants(model: "Model") -> Dict[str, Var]:
     """If there are constants available, return them as a dictionary."""
     # The constant data vars must be either pm.Data or TensorConstant or SharedVariable
-    if model is None:
-        return {}
-
     def is_data(name, var, model) -> bool:
         observations = find_observations(model)
         return (

pymc/tests/test_sampling_jax.py

@@ -171,15 +171,18 @@ def test_get_jaxified_logp():
 def test_idata_kwargs(sampler, idata_kwargs, postprocessing_backend):
     with pm.Model() as m:
         x = pm.Normal("x")
-        z = pm.Normal("z")
         y = pm.Normal("y", x, observed=0)
+        pm.ConstantData("constantdata", [1, 2, 3])
+        pm.MutableData("mutabledata", 2)
         idata = sampler(
             tune=50,
             draws=50,
             chains=1,
             idata_kwargs=idata_kwargs,
             postprocessing_backend=postprocessing_backend,
         )
+    assert "constantdata" in idata.constant_data
+    assert "mutabledata" in idata.constant_data
 
     if idata_kwargs.get("log_likelihood", True):
         assert "log_likelihood" in idata