Added alternative scale parameterization with unit tests to exponential

Author: manulpatel

pymc/distributions/continuous.py

@@ -1347,13 +1347,22 @@ class Exponential(PositiveContinuous):
     ----------
     lam : tensor_like of float
         Rate or inverse scale (``lam`` > 0).
+    scale: tensor_like of float
+        Alternative parameter (scale = 1/lam).
     """
     rv_op = exponential
 
     @classmethod
-    def dist(cls, lam: DIST_PARAMETER_TYPES, *args, **kwargs):
-        lam = pt.as_tensor_variable(floatX(lam))
+    def dist(cls, lam=None, scale=None, *args, **kwargs):
+        if lam is not None and scale is not None:
+            raise ValueError("Incompatible parametrization. Can't specify both lam and scale.")
+        elif lam is None and scale is None:
+            raise ValueError("Incompatible parametrization. Must specify either lam or scale.")
+
+        if scale is not None:
+            lam = pt.reciprocal(scale)
 
+        lam = pt.as_tensor_variable(floatX(lam))
         # PyTensor exponential op is parametrized in terms of mu (1/lam)
         return super().dist([pt.reciprocal(lam)], **kwargs)
 

tests/distributions/test_continuous.py

@@ -432,18 +432,43 @@ def test_exponential(self):
             {"lam": Rplus},
             lambda value, lam: st.expon.logpdf(value, 0, 1 / lam),
         )
+        check_logp(
+            pm.Exponential,
+            Rplus,
+            {"scale": Rplus},
+            lambda value, scale: st.expon.logpdf(value, 0, scale),
+        )
         check_logcdf(
             pm.Exponential,
             Rplus,
             {"lam": Rplus},
             lambda value, lam: st.expon.logcdf(value, 0, 1 / lam),
         )
+        check_logcdf(
+            pm.Exponential,
+            Rplus,
+            {"scale": Rplus},
+            lambda value, scale: st.expon.logcdf(value, 0, scale),
+        )
         check_icdf(
             pm.Exponential,
             {"lam": Rplus},
             lambda q, lam: st.expon.ppf(q, loc=0, scale=1 / lam),
         )
 
+    def test_exponential_wrong_arguments(self):
+        m = pm.Model()
+
+        msg = "Incompatible parametrization. Can't specify both lam and scale"
+        with m:
+            with pytest.raises(ValueError, match=msg):
+                pm.Exponential("x", lam=0.5, scale=5)
+
+        msg = "Incompatible parametrization. Must specify either lam or scale"
+        with m:
+            with pytest.raises(ValueError, match=msg):
+                pm.Exponential("x")
+
     def test_laplace(self):
         check_logp(
             pm.Laplace,
@@ -2091,6 +2116,13 @@ class TestExponential(BaseTestDistributionRandom):
     ]
 
 
+class TestExponentialScale(BaseTestDistributionRandom):
+    pymc_dist = pm.Exponential
+    pymc_dist_params = {"scale": 5.0}
+    expected_rv_op_params = {"mu": pymc_dist_params["scale"]}
+    checks_to_run = ["check_pymc_params_match_rv_op"]
+
+
 class TestCauchy(BaseTestDistributionRandom):
     pymc_dist = pm.Cauchy
     pymc_dist_params = {"alpha": 2.0, "beta": 5.0}