update to pymc 5 and pytensor (#29)

* update to pymc 5 and pytensor

* fix pylint

* black

* fix shape

Author: aloctavodia

pymc_bart/bart.py

@@ -15,16 +15,14 @@
 #   limitations under the License.
 
 from multiprocessing import Manager
-import aesara.tensor as at
 import numpy as np
-
-from aeppl.logprob import _logprob
-from aesara.tensor.random.op import RandomVariable
-from aesara.tensor.var import Variable
-
 from pandas import DataFrame, Series
 
 from pymc.distributions.distribution import Distribution, _moment
+from pymc.logprob.abstract import _logprob
+import pytensor.tensor as pt
+from pytensor.tensor.random.op import RandomVariable
+
 
 from .utils import _sample_posterior
 
@@ -42,11 +40,7 @@ class BARTRV(RandomVariable):
     all_trees = None
 
     def _supp_shape_from_params(self, dist_params, rep_param_idx=1, param_shapes=None):
-        if isinstance(self.X, Variable):
-            shape = self.X.shape[0].eval()
-        else:
-            shape = self.X.shape[0]
-        return (shape,)
+        return dist_params[0].shape[:1]
 
     @classmethod
     def rng_fn(cls, rng=None, X=None, Y=None, m=None, alpha=None, split_prior=None, size=None):
@@ -145,11 +139,11 @@ def logp(self, x, *inputs):
         -------
         TensorVariable
         """
-        return at.zeros_like(x)
+        return pt.zeros_like(x)
 
     @classmethod
     def get_moment(cls, rv, size, *rv_inputs):
-        mean = at.fill(size, rv.Y.mean())
+        mean = pt.fill(size, rv.Y.mean())
         return mean
 
 

pymc_bart/pgbart.py

@@ -19,13 +19,13 @@
 
 import numpy as np
 
-from aesara import function as aesara_function
-from aesara import config
-from aesara.tensor.var import Variable
+from pytensor import function as pytensor_function
+from pytensor import config
+from pytensor.tensor.var import Variable
 
 from pymc.model import modelcontext
 from pymc.step_methods.arraystep import ArrayStepShared, Competence
-from pymc.aesaraf import inputvars, join_nonshared_inputs, make_shared_replacements
+from pymc.pytensorf import inputvars, join_nonshared_inputs, make_shared_replacements
 
 
 from pymc_bart.bart import BARTRV
@@ -656,9 +656,9 @@ def logp(point, out_vars, vars, shared):  # pylint: disable=redefined-builtin
     vars: List
         containing :class:`pymc.Distribution` for the input variables
     shared: List
-        containing :class:`aesara.tensor.Tensor` for depended shared data
+        containing :class:`pytensor.tensor.Tensor` for depended shared data
     """
     out_list, inarray0 = join_nonshared_inputs(point, out_vars, vars, shared)
-    function = aesara_function([inarray0], out_list[0])
+    function = pytensor_function([inarray0], out_list[0])
     function.trust_input = True
     return function

pymc_bart/tree.py

@@ -16,7 +16,7 @@
 
 from copy import deepcopy
 
-from aesara import config
+from pytensor import config
 import numpy as np
 
 

pymc_bart/utils.py

@@ -4,7 +4,7 @@
 import matplotlib.pyplot as plt
 import numpy as np
 
-from aesara.tensor.var import Variable
+from pytensor.tensor.var import Variable
 from scipy.interpolate import griddata
 from scipy.signal import savgol_filter
 from scipy.stats import pearsonr