add gamma helper function to deprecate rate parameterization

Author: tvwenger

pytensor/tensor/random/basic.py

@@ -16,6 +16,7 @@
     RandomGeneratorSharedVariable,
     RandomStateSharedVariable,
 )
+import warnings
 
 
 try:
@@ -443,7 +444,7 @@ class GammaRV(ScipyRandomVariable):
     dtype = "floatX"
     _print_name = ("Gamma", "\\operatorname{Gamma}")
 
-    def __call__(self, alpha, beta, size=None, **kwargs):
+    def __call__(self, shape, scale, size=None, **kwargs):
         r"""Draw samples from a gamma distribution.
 
         Signature
@@ -455,23 +456,39 @@ def __call__(self, alpha, beta, size=None, **kwargs):
         ----------
         shape
             The shape :math:`\alpha` of the gamma distribution. Must be positive.
-        rate
-            The rate :math:`\beta` of the gamma distribution. Must be positive.
+        scale
+            The scale :math:`1/\beta` of the gamma distribution. Must be positive.
         size
             Sample shape. If the given size is, e.g. `(m, n, k)` then `m * n * k`
             independent, identically distributed random variables are
             returned. Default is `None` in which case a single random variable
             is returned.
 
         """
-        return super().__call__(alpha, beta, size=size, **kwargs)
+        return super().__call__(shape, scale, size=size, **kwargs)
 
     @classmethod
-    def rng_fn_scipy(cls, rng, alpha, beta, size):
-        return stats.gamma.rvs(alpha, scale=1.0 / beta, size=size, random_state=rng)
+    def rng_fn_scipy(cls, rng, shape, scale, size):
+        return stats.gamma.rvs(shape, scale=scale, size=size, random_state=rng)
+
+
+_gamma = GammaRV()
 
 
-gamma = GammaRV()
+def gamma(shape, scale=None, rate=None, **kwargs):
+    # TODO: Remove helper when rate is deprecated
+    if rate is not None and scale is not None:
+        raise ValueError("Can't specify both rate and scale")
+    elif rate is None and scale is None:
+        raise ValueError("Must specify scale")
+    elif rate is not None:
+        warnings.warn(
+            "Gamma rate argument is deprecated and will stop working, use scale instead",
+            FutureWarning,
+        )
+        scale = 1.0 / rate
+
+    return _gamma(shape, scale, **kwargs)
 
 
 class ChiSquareRV(RandomVariable):

tests/link/jax/test_random.py

@@ -147,8 +147,8 @@ def test_random_updates_input_storage_order():
                     np.array([1.0, 2.0], dtype=np.float64),
                 ),
                 set_test_value(
-                    at.dscalar(),
-                    np.array(1.0, dtype=np.float64),
+                    at.dvector(),
+                    np.array([0.5, 3.0], dtype=np.float64),
                 ),
             ],
             (2,),

tests/link/numba/test_random.py

@@ -323,8 +323,8 @@ def test_aligned_RandomVariable(rv_op, dist_args, size):
                     np.array([1.0, 2.0], dtype=np.float64),
                 ),
                 set_test_value(
-                    at.dscalar(),
-                    np.array(1.0, dtype=np.float64),
+                    at.dvector(),
+                    np.array([0.5, 3.0], dtype=np.float64),
                 ),
             ],
             (2,),

tests/tensor/random/test_basic.py

@@ -353,8 +353,8 @@ def test_lognormal_samples(mean, sigma, size):
 def test_gamma_samples(a, b, size):
     gamma_test_fn = fixed_scipy_rvs("gamma")
 
-    def test_fn(shape, rate, **kwargs):
-        return gamma_test_fn(shape, scale=1.0 / rate, **kwargs)
+    def test_fn(shape, scale, **kwargs):
+        return gamma_test_fn(shape, scale=scale, **kwargs)
 
     compare_sample_values(
         gamma,