fix split_prior bug (#115)

* fix split_prior bug

* fix type

* fix type

Author: aloctavodia

pymc_bart/bart.py

@@ -43,7 +43,9 @@ class BARTRV(RandomVariable):
     _print_name: Tuple[str, str] = ("BART", "\\operatorname{BART}")
     all_trees = List[List[List[Tree]]]
 
-    def _supp_shape_from_params(self, dist_params, rep_param_idx=1, param_shapes=None):
+    def _supp_shape_from_params(
+        self, dist_params, rep_param_idx=1, param_shapes=None
+    ):  # pylint: disable=arguments-renamed
         return dist_params[0].shape[:1]
 
     @classmethod
@@ -126,7 +128,7 @@ def __new__(
         alpha: float = 0.95,
         beta: float = 2.0,
         response: str = "constant",
-        split_prior: Optional[List[float]] = None,
+        split_prior: Optional[npt.NDArray[np.float_]] = None,
         split_rules: Optional[List[SplitRule]] = None,
         separate_trees: Optional[bool] = False,
         **kwargs,
@@ -141,8 +143,7 @@ def __new__(
 
         X, Y = preprocess_xy(X, Y)
 
-        if split_prior is None:
-            split_prior = []
+        split_prior = np.array([]) if split_prior is None else np.asarray(split_prior)
 
         bart_op = type(
             f"BART_{name}",

pymc_bart/pgbart.py

@@ -145,10 +145,10 @@ def __init__(
         self.trees_shape = self.shape if self.bart.separate_trees else 1
         self.leaves_shape = self.shape if not self.bart.separate_trees else 1
 
-        if self.bart.split_prior:
-            self.alpha_vec = self.bart.split_prior
+        if self.bart.split_prior.size == 0:
+            self.alpha_vec = np.ones(self.X.shape[1])
         else:
-            self.alpha_vec = np.ones(self.X.shape[1], dtype=np.int32)
+            self.alpha_vec = self.bart.split_prior
 
         if self.bart.split_rules:
             self.split_rules = self.bart.split_rules