Move chi_dist into class

Author: ricardoV94

pymc_experimental/distributions/continuous.py

@@ -220,10 +220,6 @@ def moment(rv, size, mu, sigma, xi):
         return mode
 
 
-def chi_dist(nu: TensorVariable, size: TensorVariable) -> TensorVariable:
-    return pt.math.sqrt(ChiSquared.dist(nu=nu, size=size))
-
-
 class Chi:
     r"""
     :math:`\chi` log-likelihood.
@@ -270,11 +266,15 @@ class Chi:
             x = pm.Chi('x', nu=1)
     """
 
+    @staticmethod
+    def chi_dist(nu: TensorVariable, size: TensorVariable) -> TensorVariable:
+        return pt.math.sqrt(ChiSquared.dist(nu=nu, size=size))
+
     def __new__(cls, name, nu, **kwargs):
         if "observed" not in kwargs:
             kwargs.setdefault("transform", transforms.log)
-        return CustomDist(name, nu, dist=chi_dist, class_name="Chi", **kwargs)
+        return CustomDist(name, nu, dist=cls.chi_dist, class_name="Chi", **kwargs)
 
     @classmethod
     def dist(cls, nu, **kwargs):
-        return CustomDist.dist(nu, dist=chi_dist, class_name="Chi", **kwargs)
+        return CustomDist.dist(nu, dist=cls.chi_dist, class_name="Chi", **kwargs)