Reparameterize GammaRV so beta is not inverted at each call

Also fix wrong JAX implementation of Gamma and Pareto RVs

Author: tvwenger

pytensor/link/jax/dispatch/random.py

@@ -214,23 +214,22 @@ def sample_fn(rng, size, dtype, *parameters):
     return sample_fn
 
 
-@jax_sample_fn.register(aer.ParetoRV)
 @jax_sample_fn.register(aer.GammaRV)
-def jax_sample_fn_shape_rate(op):
-    """JAX implementation of random variables in the shape-rate family.
+@jax_sample_fn.register(aer.ParetoRV)
+def jax_sample_fn_shape_scale(op):
+    """JAX implementation of random variables in the shape-scale family.
 
     JAX only implements the standard version of random variables in the
-    shape-rate family. We thus need to rescale the results manually.
+    shape-scale family. We thus need to rescale the results manually.
 
     """
     name = op.name
     jax_op = getattr(jax.random, name)
 
-    def sample_fn(rng, size, dtype, *parameters):
+    def sample_fn(rng, size, dtype, shape, scale):
         rng_key = rng["jax_state"]
         rng_key, sampling_key = jax.random.split(rng_key, 2)
-        (shape, rate) = parameters
-        sample = jax_op(sampling_key, shape, size, dtype) / rate
+        sample = jax_op(sampling_key, shape, size, dtype) * scale
         rng["jax_state"] = rng_key
         return (rng, sample)
 

pytensor/tensor/random/basic.py

@@ -1,4 +1,5 @@
 import abc
+import warnings
 from typing import List, Optional, Union
 
 import numpy as np
@@ -419,7 +420,7 @@ def __call__(self, mean=0.0, sigma=1.0, size=None, **kwargs):
 lognormal = LogNormalRV()
 
 
-class GammaRV(ScipyRandomVariable):
+class GammaRV(RandomVariable):
     r"""A gamma continuous random variable.
 
     The probability density function for `gamma` in terms of the shape parameter
@@ -443,7 +444,7 @@ class GammaRV(ScipyRandomVariable):
     dtype = "floatX"
     _print_name = ("Gamma", "\\operatorname{Gamma}")
 
-    def __call__(self, shape, rate, 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, shape, rate, 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__(shape, 1.0 / rate, size=size, **kwargs)
+        return super().__call__(shape, scale, size=size, **kwargs)
 
     @classmethod
     def rng_fn_scipy(cls, rng, shape, scale, size):
-        return stats.gamma.rvs(shape, scale=scale, size=size, random_state=rng)
+        return rng.gamma(shape, scale=scale, size=size)
+
 
+_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("Cannot 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
 
-gamma = GammaRV()
+    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,),
@@ -235,11 +235,15 @@ def test_random_updates_input_storage_order():
                 set_test_value(
                     at.dvector(),
                     np.array([1.0, 2.0], dtype=np.float64),
-                )
+                ),
+                set_test_value(
+                    at.dvector(),
+                    np.array([2.0, 10.0], dtype=np.float64),
+                ),
             ],
             (2,),
             "pareto",
-            lambda *args: args,
+            lambda shape, scale: (shape, 0.0, scale),
         ),
         (
             aer.poisson,

tests/link/numba/test_random.py

@@ -92,6 +92,10 @@
                     at.dvector(),
                     np.array([1.0, 2.0], dtype=np.float64),
                 ),
+                set_test_value(
+                    at.dvector(),
+                    np.array([2.0, 10.0], dtype=np.float64),
+                ),
             ],
             at.as_tensor([3, 2]),
             marks=pytest.mark.xfail(reason="Not implemented"),
@@ -323,8 +327,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

@@ -351,17 +351,11 @@ 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)
-
     compare_sample_values(
         gamma,
         a,
         b,
         size=size,
-        test_fn=test_fn,
     )
 
 
@@ -470,18 +464,24 @@ def test_vonmises_samples(mu, kappa, size):
 
 
 @pytest.mark.parametrize(
-    "alpha, size",
+    "alpha, scale, size",
     [
-        (np.array(0.5, dtype=config.floatX), None),
-        (np.array(0.5, dtype=config.floatX), []),
+        (np.array(0.5, dtype=config.floatX), np.array(3.0, dtype=config.floatX), None),
+        (np.array(0.5, dtype=config.floatX), np.array(5.0, dtype=config.floatX), []),
         (
             np.full((1, 2), 0.5, dtype=config.floatX),
+            np.array([0.5, 1.0], dtype=config.floatX),
             None,
         ),
     ],
 )
-def test_pareto_samples(alpha, size):
-    compare_sample_values(pareto, alpha, size=size, test_fn=fixed_scipy_rvs("pareto"))
+def test_pareto_samples(alpha, scale, size):
+    pareto_test_fn = fixed_scipy_rvs("pareto")
+
+    def test_fn(shape, scale, **kwargs):
+        return pareto_test_fn(shape, scale=scale, **kwargs)
+
+    compare_sample_values(pareto, alpha, scale, size=size, test_fn=test_fn)
 
 
 def mvnormal_test_fn(mean=None, cov=None, size=None, random_state=None):