add constant_data into sample_numpyro_nuts()

Author: danhphan

pymc/backends/arviz.py

@@ -63,6 +63,22 @@ def find_observations(model: Optional["Model"]) -> Optional[Dict[str, Var]]:
     return observations
 
 
+def find_constants(model: Optional["Model"]) -> Optional[Dict[str, Var]]:
+    """If there are constants available, return them as a dictionary."""
+    if model is None or not model.named_vars:
+        return None
+
+    constants = {}
+    for name, var in model.named_vars.items():
+        if isinstance(var, (Constant, SharedVariable)):
+            if hasattr(var, "data"):
+                var = var.data
+            elif hasattr(var, "get_value"):
+                var = var.get_value()
+            constants[name] = var
+    return constants
+
+
 class _DefaultTrace:
     """
     Utility for collecting samples into a dictionary.

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.")
@@ -564,6 +564,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,