Prior: Help users transition from non-scalar parameters to DataArray

Author: ricardoV94

pymc_extras/prior.py

@@ -85,6 +85,7 @@ def custom_transform(x):
 
 import copy
 import typing
+import warnings
 
 from collections.abc import Callable
 from functools import partial
@@ -493,6 +494,39 @@ def create_variable(self, name: str) -> "TensorVariable":
     ).prior
 
 
+def _param_value_with_dims(param, value, dims: Dims | None):
+    """Infer parameter dims positionally.
+
+    This is a transition helper to guide users into defining DataArray parameters explicitly.
+    """
+    if hasattr(value, "dims"):
+        return value
+
+    if isinstance(value, list | tuple | Number):
+        value = np.asarray(value)
+
+    if value.ndim > 0:
+        if dims is None:
+            raise ValueError(
+                f"Cannot infer dims of array-like parameter {param}. Use DataArray with explicit dims"
+            )
+        else:
+            parameter_dims = dims[::-1][: value.ndim]
+            warnings.warn(
+                f"Implicit conversion of array-like parameter {param} to DataArray with dims {parameter_dims}. "
+                "Use DataArray with explicit dims to avoid this warning",
+                stacklevel=2,
+            )
+            if isinstance(value, Variable):
+                from pytensor.xtensor import as_xtensor
+
+                value = as_xtensor(value, dims=parameter_dims)
+            else:
+                value = DataArray(value, dims=parameter_dims)
+
+    return value
+
+
 class Prior:
     """A class to represent a prior distribution.
 
@@ -777,7 +811,10 @@ def __repr__(self) -> str:
 
     def _create_parameter(self, param, value, name, xdist: bool = False):
         if not hasattr(value, "create_variable"):
-            return value
+            if xdist:
+                return _param_value_with_dims(param, value, dims=self.dims)
+            else:
+                return value
 
         child_name = f"{name}_{param}"
         if xdist:
@@ -803,7 +840,10 @@ def _create_non_centered_variable(
         def handle_variable(var_name: str):
             parameter = self.parameters[var_name]
             if not hasattr(parameter, "create_variable"):
-                return parameter
+                if xdist:
+                    return _param_value_with_dims(var_name, parameter, dims=self.dims)
+                else:
+                    return parameter
 
             if xdist:
                 return parameter.create_variable(f"{name}_{var_name}", xdist=True)

tests/test_prior.py

@@ -1327,3 +1327,25 @@ def test_dims(self) -> None:
         # This is always invalid
         with pytest.raises(UnsupportedShapeError):
             p.dims = ()
+
+    @pytest.mark.parametrize(
+        "mu",
+        (
+            [10, 20],
+            np.array([10, 20]),
+            pt.as_tensor([10, 20]),
+        ),
+    )
+    def test_implicit_conversion_with_dims(self, mu):
+        # When xdist=True, list/numpy parameters should be converted to DataArray if dims are present
+        p_wo_dims = Prior("Normal", mu=mu, dims=None)
+        p_with_dims = Prior("Normal", mu=mu, dims=("test_dim",))
+
+        coords = {"test_dim": ["a", "b"]}
+        with pm.Model(coords=coords):
+            with pytest.warns(UserWarning, match="Implicit conversion"):
+                res = p_with_dims.create_variable("v", xdist=True)
+                assert res.dims == ("test_dim",)
+
+            with pytest.raises(ValueError, match="Cannot infer dims"):
+                p_wo_dims.create_variable("v", xdist=True)