Updated ChiSquare r.v. and added tests

aesara/tensor/random/basic.py

@@ -117,32 +117,15 @@ class GammaRV(RandomVariable):
 gamma = GammaRV()
 
 
-class ChiSquaredRV(RandomVariable):
-    name = "chi_squared"
+class ChiSquareRV(RandomVariable):
+    name = "chisquare"
     ndim_supp = 0
     ndims_params = [0]
     dtype = "floatX"
-    _print_name = ("ChiSquared", "\\operatorname{ChiSquared}")
-
-    def __call__(
-        self,
-        nu: Union[np.ndarray, float],
-        size: Optional[Union[List[int], int]] = None,
-        **kwargs
-    ) -> RandomVariable:
-        return super().call(nu=nu, size=size, **kwargs)
-
-    @classmethod
-    def rng_fn(
-        cls,
-        rng: np.random.RandomState,
-        nu: Union[np.ndarray, float],
-        size: Optional[Union[List[int], int]],
-    ) -> np.ndarray:
-        return stats.chi2.rvs(nu, size=size, random_state=rng)
+    _print_name = ("ChiSquare", "\\operatorname{ChiSquare}")
 
 
-chi_squared = ChiSquaredRV()
+chisquare = ChiSquareRV()
 
 
 class ParetoRV(RandomVariable):

tests/tensor/random/test_basic.py

@@ -19,6 +19,7 @@ from aesara.tensor.random.basic import (
     binomial,
     categorical,
     cauchy,
+    chisquare,
     choice,
     dirichlet,
     exponential,
@@ -244,6 +245,13 @@ def test_gamma_samples():
     rv_numpy_tester(gamma, test_a, test_b, size=[2, 3], test_fn=stats.gamma.rvs)
 
 
+def test_chisquare_samples():
+    test_df = np.array(2, dtype=config.floatX)
+
+    rv_numpy_tester(chisquare, test_df, test_fn=stats.chi2.rvs)
+    rv_numpy_tester(chisquare, test_df, size=[2, 3], test_fn=stats.gamma.rvs)
+
+
 def test_gumbel_samples():
     test_mu = np.array(0.0, dtype=config.floatX)
     test_beta = np.array(1.0, dtype=config.floatX)