Enforce check_parameters for alpha in geometric_adstock

pymc_marketing/mmm/transformers.py

@@ -20,6 +20,7 @@
 import numpy.typing as npt
 import pymc as pm
 import pytensor.tensor as pt
+from pymc.distributions.dist_math import check_parameters
 from pytensor.tensor.random.utils import params_broadcast_shapes
 
 
@@ -235,6 +236,10 @@ def geometric_adstock(
        with carryover and shape effects." (2017).
 
     """
+    alpha = check_parameters(
+        alpha, [pt.ge(alpha, 0), pt.le(alpha, 1)], msg="0 <= alpha <= 1"
+    )
+
     w = pt.power(pt.as_tensor(alpha)[..., None], pt.arange(l_max, dtype=x.dtype))
     w = w / pt.sum(w, axis=-1, keepdims=True) if normalize else w
     return batched_convolution(x, w, axis=axis, mode=mode)

tests/mmm/test_transformers.py

@@ -18,6 +18,7 @@
 import pytensor.tensor as pt
 import pytest
 import scipy as sp
+from pymc.logprob.utils import ParameterValueError
 from pytensor.tensor.variable import TensorVariable
 
 from pymc_marketing.mmm.transformers import (
@@ -148,6 +149,21 @@ def test_geometric_adstock_good_alpha(self, x, alpha, l_max):
         assert y_np[1] == x[1] + alpha * x[0]
         assert y_np[2] == x[2] + alpha * x[1] + (alpha**2) * x[0]
 
+    @pytest.mark.parametrize(
+        "x, alpha, l_max",
+        [
+            (np.ones(shape=(100)), -0.3, 10),
+            (np.ones(shape=(100)), -1.7, 100),
+            (np.ones(shape=(100)), 22.5, 7),
+            (np.linspace(start=0.0, stop=1.0, num=50), -0.8, 3),
+            (np.linspace(start=0.0, stop=1.0, num=50), 8, 50),
+        ],
+    )
+    def test_geometric_adstock_bad_alpha(self, x, alpha, l_max):
+        with pytest.raises(ParameterValueError):
+            y = geometric_adstock(x=x, alpha=alpha, l_max=l_max)
+            y.eval()
+
     @pytest.mark.parametrize(
         argnames="mode",
         argvalues=[ConvMode.After, ConvMode.Before, ConvMode.Overlap],