Made existing tests work and added a new test for separate_trees

Author: velochy

pymc_bart/utils.py

@@ -49,7 +49,7 @@ def _sample_posterior(
         X = X.eval()
 
     if size is None:
-        size_iter: Union[List, Tuple] = ()
+        size_iter: Union[List, Tuple] = (1,)
     elif isinstance(size, int):
         size_iter = [size]
     else:

tests/test_bart.py

@@ -217,3 +217,23 @@ def test_bart_moment(size, expected):
     with pm.Model() as model:
         pmb.BART("x", X=X, Y=Y, size=size)
     assert_moment_is_expected(model, expected)
+
+
+@pytest.mark.parametrize(
+    argnames="separate_trees",
+    argvalues=[False, True],
+    ids=["shared-trees", "separate-trees"],
+)
+def test_categorical(separate_trees):
+
+    Y = np.array([0, 0, 0, 1, 1, 1, 2, 2, 2])
+    X = np.concatenate([Y[:, None], np.random.randint(0, 6, size=(9, 4))], axis=1)
+
+    with pm.Model() as model:
+        lo = pmb.BART("logodds", X, Y, m=2, shape=(3, 9), 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)
+
+    # Fit should be good enough so right category is selected over 50% of time
+    assert (idata.predictions.y.median(["chain", "draw"]) == Y).all()