Add the standard gamma RandomVariable

Author: rlouf

aesara/tensor/random/basic.py

@@ -431,6 +431,45 @@ def rng_fn_scipy(cls, rng, shape, scale, size):
 gamma = GammaRV()
 
 
+class StandardGammaRV(GammaRV):
+    r"""A standard gamma continuous random variable.
+
+    The probability density function for `standard_gamma` in terms of its shape
+    parameters :math:`\alpha` is:
+
+    .. math::
+
+        f(x) = \frac{1}{\Gamma(\alpha)}x^{\alpha-1}e^{-x}
+
+    for :math:`x \geq 0`, :math:`\alpha > 0`. :math:`\Gamma` is the gamma
+    function:
+
+    .. math::
+
+        \Gamma(x) = \int_0^{\infty} t^{x-1} e^{-t} \mathrm{d}t
+
+    """
+
+    def __call__(self, shape, size=None, **kwargs):
+        """Draw samples from a standard gamma distribution.
+
+        Parameters
+        ----------
+        shape
+            The shape :math:`\alpha` 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, rate=1.0, size=size, **kwargs)
+
+
+standard_gamma = StandardGammaRV()
+
+
 class ChiSquareRV(RandomVariable):
     r"""A chi square continuous random variable.
 
@@ -2012,6 +2051,7 @@ def __call__(self, x, **kwargs):
     "uniform",
     "standard_cauchy",
     "standard_exponential",
+    "standard_gamma",
     "standard_normal",
     "standard_t",
     "negative_binomial",

tests/tensor/random/test_basic.py

@@ -49,6 +49,7 @@
     randint,
     standard_cauchy,
     standard_exponential,
+    standard_gamma,
     standard_normal,
     standard_t,
     t,
@@ -368,6 +369,31 @@ def test_fn(shape, rate, **kwargs):
     )
 
 
+@pytest.mark.parametrize(
+    "a, size",
+    [
+        (np.array(0.5, dtype=config.floatX), None),
+        (np.array(0.5, dtype=config.floatX), []),
+        (
+            np.full((1, 2), 0.5, dtype=config.floatX),
+            None,
+        ),
+    ],
+)
+def test_standard_gamma_samples(a, size):
+    gamma_test_fn = fixed_scipy_rvs("gamma")
+
+    def test_fn(shape, **kwargs):
+        return gamma_test_fn(shape, scale=1.0, **kwargs)
+
+    compare_sample_values(
+        standard_gamma,
+        a,
+        size=size,
+        test_fn=test_fn,
+    )
+
+
 @pytest.mark.parametrize(
     "df, size",
     [