some fixes

Author: Juan Orduz

pymc_bart/pgbart.py

@@ -556,18 +556,18 @@ def draw_leaf_value(
 ) -> tuple[npt.NDArray[np.float64], Optional[npt.NDArray[np.float64]]]:
     """Draw Gaussian distributed leaf values."""
     linear_params = None
-
+    mu_mean: npt.NDArray[np.float64] = np.empty(shape)
     if y_mu_pred.size == 0:
         return np.zeros(shape), linear_params
 
     if y_mu_pred.size == 1:
-        mu_mean = np.full(shape, y_mu_pred.item() / m) + norm
+        mu_mean = (np.full(shape, y_mu_pred.item() / m) + norm).astype(np.float64)
     elif y_mu_pred.size < 3 or response == "constant":
         mu_mean = fast_mean(y_mu_pred) / m + norm
     else:
         mu_mean, linear_params = fast_linear_fit(x=x_mu, y=y_mu_pred, m=m, norm=norm)
 
-    return (mu_mean).astype(np.float64), linear_params
+    return mu_mean, linear_params
 
 
 @njit

pymc_bart/utils.py

@@ -826,9 +826,9 @@ def compute_variable_importance(  # noqa: PLR0915 PLR0912
     else:
         labels = np.arange(n_vars).astype(str)
 
-    r2_mean = np.zeros(n_vars)
-    r2_hdi = np.zeros((n_vars, 2))
-    preds = np.zeros((n_vars, samples, *bartrv.eval().T.shape))
+    r2_mean: npt.NDArray[np.float64] = np.zeros(n_vars)
+    r2_hdi: npt.NDArray[np.float64] = np.zeros((n_vars, 2))
+    preds: npt.NDArray[np.float64] = np.zeros((n_vars, samples, *bartrv.eval().T.shape))
 
     if method == "backward_VI":
         if fixed >= n_vars: