Change to Distribution (#15)

* change to Distribution

* fix test

Co-authored-by: aloctavodia <[email protected]>

Author: PabloGGaray

pymc_bart/bart.py

@@ -21,7 +21,7 @@
 from aesara.tensor.random.op import RandomVariable
 from pandas import DataFrame, Series
 
-from pymc.distributions.distribution import NoDistribution, _moment
+from pymc.distributions.distribution import Distribution, _moment
 
 __all__ = ["BART"]
 
@@ -50,7 +50,7 @@ def rng_fn(cls, rng, X, Y, m, alpha, split_prior, size):
 bart = BARTRV()
 
 
-class BART(NoDistribution):
+class BART(Distribution):
     """
     Bayesian Additive Regression Tree distribution.
 
@@ -104,7 +104,7 @@ def __new__(
             ),
         )()
 
-        NoDistribution.register(BARTRV)
+        Distribution.register(BARTRV)
 
         @_moment.register(BARTRV)
         def get_moment(rv, size, *rv_inputs):

tests/test_bart.py

@@ -3,10 +3,11 @@
 import pytest
 from numpy.random import RandomState
 from numpy.testing import assert_almost_equal, assert_array_equal
-from pymc.tests.test_distributions_moments import assert_moment_is_expected
+from pymc.tests.distributions.util import assert_moment_is_expected
 
 import pymc_bart as pmb
 
+
 def test_split_node():
     split_node = pmb.tree.SplitNode(index=5, idx_split_variable=2, split_value=3.0)
     assert split_node.index == 5