Add constant_data into sample_numpyro_nuts() (#5807)

Add constant_data into sample_numpyro_nuts()

Co-authored-by: Michael Osthege <[email protected]>

Author: danhphan

pymc/backends/arviz.py

@@ -43,11 +43,8 @@
 Var = Any  # pylint: disable=invalid-name
 
 
-def find_observations(model: Optional["Model"]) -> Optional[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 None
-
     observations = {}
     for obs in model.observed_RVs:
         aux_obs = getattr(obs.tag, "observations", None)
@@ -63,6 +60,36 @@ def find_observations(model: Optional["Model"]) -> Optional[Dict[str, Var]]:
     return observations
 
 
+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
+    def is_data(name, var, model) -> bool:
+        observations = find_observations(model)
+        return (
+            var not in model.deterministics
+            and var not in model.observed_RVs
+            and var not in model.free_RVs
+            and var not in model.potentials
+            and var not in model.value_vars
+            and name not in observations
+            and isinstance(var, (Constant, SharedVariable))
+        )
+
+    # The assumption is that constants (like pm.Data) are named
+    # variables that aren't observed or free RVs, nor are they
+    # deterministics, and then we eliminate observations.
+    constant_data = {}
+    for name, var in model.named_vars.items():
+        if is_data(name, var, model):
+            if hasattr(var, "get_value"):
+                var = var.get_value()
+            elif hasattr(var, "data"):
+                var = var.data
+            constant_data[name] = var
+
+    return constant_data
+
+
 class _DefaultTrace:
     """
     Utility for collecting samples into a dictionary.
@@ -467,41 +494,10 @@ def observed_data_to_xarray(self):
     @requires("model")
     def constant_data_to_xarray(self):
         """Convert constant data to xarray."""
-        # For constant data, we are concerned only with deterministics and
-        # data.  The constant data vars must be either pm.Data
-        # (TensorConstant/SharedVariable) or pm.Deterministic
-        constant_data_vars = {}  # type: Dict[str, Var]
-
-        def is_data(name, var) -> bool:
-            assert self.model is not None
-            return (
-                var not in self.model.deterministics
-                and var not in self.model.observed_RVs
-                and var not in self.model.free_RVs
-                and var not in self.model.potentials
-                and var not in self.model.value_vars
-                and (self.observations is None or name not in self.observations)
-                and isinstance(var, (Constant, SharedVariable))
-            )
-
-        # I don't know how to find pm.Data, except that they are named
-        # variables that aren't observed or free RVs, nor are they
-        # deterministics, and then we eliminate observations.
-        for name, var in self.model.named_vars.items():
-            if is_data(name, var):
-                constant_data_vars[name] = var
-
-        if not constant_data_vars:
+        constant_data = find_constants(self.model)
+        if not constant_data:
             return None
 
-        constant_data = {}
-        for name, vals in constant_data_vars.items():
-            if hasattr(vals, "get_value"):
-                vals = vals.get_value()
-            elif hasattr(vals, "data"):
-                vals = vals.data
-            constant_data[name] = vals
-
         return dict_to_dataset(
             constant_data,
             library=pymc,

pymc/sampling_jax.py

@@ -30,7 +30,7 @@
 from arviz.data.base import make_attrs
 
 from pymc import Model, modelcontext
-from pymc.backends.arviz import find_observations
+from pymc.backends.arviz import find_constants, find_observations
 from pymc.util import get_default_varnames
 
 warnings.warn("This module is experimental.")
@@ -370,6 +370,7 @@ def sample_blackjax_nuts(
         posterior=posterior,
         log_likelihood=log_likelihood,
         observed_data=find_observations(model),
+        constant_data=find_constants(model),
         coords=coords,
         dims=dims,
         attrs=make_attrs(attrs, library=blackjax),
@@ -564,6 +565,7 @@ def sample_numpyro_nuts(
         posterior=posterior,
         log_likelihood=log_likelihood,
         observed_data=find_observations(model),
+        constant_data=find_constants(model),
         sample_stats=_sample_stats_to_xarray(pmap_numpyro),
         coords=coords,
         dims=dims,

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