pull out function get_constant_coords fn and add test

Author: JasonTam

pymc/sampling/forward.py

@@ -76,6 +76,28 @@
 _log = logging.getLogger(__name__)
 
 
+def get_constant_coords(trace_coords: dict[str, np.ndarray], model: Model) -> set:
+    """Get the set of coords that have remained constant between the trace and model"""
+    constant_coords = set()
+    for dim, coord in trace_coords.items():
+        current_coord = model.coords.get(dim, None)
+        current_length = model.dim_lengths.get(dim, None)
+        if hasattr(current_length, "get_value"):
+            current_length = current_length.get_value()
+        elif hasattr(current_length, "data"):
+            current_length = current_length.data
+        if (
+            current_coord is not None
+            and len(coord) == len(current_coord)
+            and np.all(coord == current_coord)
+        ) or (
+            # Coord was defined without values (only length)
+            current_coord is None and len(coord) == current_length
+        ):
+            constant_coords.add(dim)
+    return constant_coords
+
+
 def get_vars_in_point_list(trace, model):
     """Get the list of Variable instances in the model that have values stored in the trace."""
     if not isinstance(trace, MultiTrace):
@@ -792,22 +814,7 @@ def sample_posterior_predictive(
             stacklevel=2,
         )
 
-    constant_coords = set()
-    for dim, coord in trace_coords.items():
-        current_coord = model.coords.get(dim, None)
-        current_length = model.dim_lengths.get(dim, None)
-        if hasattr(current_length, 'eval'):
-            current_length = current_length.eval()
-        if (
-            current_coord is not None
-            and len(coord) == len(current_coord)
-            and np.all(coord == current_coord)
-        ) or (
-            # Coord was defined without values (only length)
-            current_coord is None
-            and len(coord) == current_length
-        ):
-            constant_coords.add(dim)
+    constant_coords = get_constant_coords(trace_coords, model)
 
     if var_names is not None:
         vars_ = [model[x] for x in var_names]

tests/sampling/test_forward.py

@@ -35,6 +35,7 @@
 from pymc.pytensorf import compile_pymc
 from pymc.sampling.forward import (
     compile_forward_sampling_function,
+    get_constant_coords,
     get_vars_in_point_list,
     observed_dependent_deterministics,
 )
@@ -1669,6 +1670,17 @@ def test_Triangular(
         assert prior["target"].shape == (prior_samples, *shape)
 
 
+def test_get_constant_coords():
+    with pm.Model() as model:
+        model.add_coord("coord0", length=1)
+
+    trace_coords_same_len = {"coord0": np.array([0])}
+    assert "coord0" in get_constant_coords(trace_coords_same_len, model)
+
+    trace_coords_diff_len = {"coord0": np.array([0, 1])}
+    assert "coord0" not in get_constant_coords(trace_coords_diff_len, model)
+
+
 def test_get_vars_in_point_list():
     with pm.Model() as modelA:
         pm.Normal("a", 0, 1)