fix zip and new args

Author: aloctavodia

pymc_bart/bart.py

@@ -175,7 +175,7 @@ def get_moment(rv, size, *rv_inputs):
             return cls.get_moment(rv, size, *rv_inputs)
 
         cls.rv_op = bart_op
-        params = [X, Y, m, alpha, beta, split_prior]
+        params = [X, Y, m, alpha, beta]
         return super().__new__(cls, name, *params, **kwargs)
 
     @classmethod

pymc_bart/pgbart.py

@@ -17,6 +17,7 @@
 import numpy as np
 import numpy.typing as npt
 from numba import njit
+from pymc.initial_point import PointType
 from pymc.model import Model, modelcontext
 from pymc.pytensorf import inputvars, join_nonshared_inputs, make_shared_replacements
 from pymc.step_methods.arraystep import ArrayStepShared
@@ -125,9 +126,12 @@ def __init__(  # noqa: PLR0915
         num_particles: int = 10,
         batch: tuple[float, float] = (0.1, 0.1),
         model: Optional[Model] = None,
+        initial_point: PointType | None = None,
+        compile_kwargs: dict | None = None,  # pylint: disable=unused-argument
     ):
         model = modelcontext(model)
-        initial_values = model.initial_point()
+        if initial_point is None:
+            initial_point = model.initial_point()
         if vars is None:
             vars = model.value_vars
         else:
@@ -150,7 +154,7 @@ def __init__(  # noqa: PLR0915
         self.m = self.bart.m
         self.response = self.bart.response
 
-        shape = initial_values[value_bart.name].shape
+        shape = initial_point[value_bart.name].shape
 
         self.shape = 1 if len(shape) == 1 else shape[0]
 
@@ -217,8 +221,8 @@ def __init__(  # noqa: PLR0915
 
         self.num_particles = num_particles
         self.indices = list(range(1, num_particles))
-        shared = make_shared_replacements(initial_values, vars, model)
-        self.likelihood_logp = logp(initial_values, [model.datalogp], vars, shared)
+        shared = make_shared_replacements(initial_point, vars, model)
+        self.likelihood_logp = logp(initial_point, [model.datalogp], vars, shared)
         self.all_particles = [
             [ParticleTree(self.a_tree) for _ in range(self.m)] for _ in range(self.trees_shape)
         ]

tests/test_bart.py

@@ -248,8 +248,10 @@ def test_categorical_model(separate_trees, split_rule):
             separate_trees=separate_trees,
         )
         y = pm.Categorical("y", p=pm.math.softmax(lo.T, axis=-1), observed=Y)
-        idata = pm.sample(random_seed=3415, tune=300, draws=300)
-        idata = pm.sample_posterior_predictive(idata, predictions=True, extend_inferencedata=True)
+        idata = pm.sample(tune=300, draws=300, random_seed=3415)
+        idata = pm.sample_posterior_predictive(
+            idata, predictions=True, extend_inferencedata=True, random_seed=3415
+        )
 
     # Fit should be good enough so right category is selected over 50% of time
     assert (idata.predictions.y.median(["chain", "draw"]) == Y).all()