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,29 @@ 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([
+        [1,1,1,0,0,0,0,0,0],
+        [0,0,0,1,1,1,0,0,0],
+        [0,0,0,0,0,0,1,1,1],
+    ]).T
+
+    X = np.concatenate([Y[:,:2], np.random.randint(0,6,size=(9,4))],axis=1)
+
+    with pm.Model() as model:
+        w = pmb.BART("w", X, Y, m=3, shape=Y.T.shape, separate_trees=separate_trees)
+        y = pm.Multinomial("y", p=pm.math.softmax(w.T,axis=-1), n=1, observed=Y, shape=Y.shape)
+        idata = pm.sample(random_seed=3415, tune=300, draws=300)
+        idata = pm.sample_posterior_predictive(idata, predictions=True, extend_inferencedata=True)
+
+    probs = np.array(idata.predictions.y.mean(['chain','draw']))[Y.astype(bool)] 
+
+    assert (probs>0.4).all()