Hello,

I’ve been looking into various use-cases for BART models and in particular have read through the Ch7 example in Bayesian Modeling and Computation in Python as well as the similar example notebook. I also looked at this binary classification example.

I was having trouble getting a binary classification model to work because I was unable to recover the labels after `pmx.bart.predict`

for an out-of-sample (hold-out) data set.

My data consists of five (5) variables and a binary response variable.

```
import pymc as pm
import pymc_experimental as pmx
from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y_vals, test_size=0.25, random_state=42, stratify= y_vals)
with pm.Model() as bart_model:
μ = pmx.BART("μ", X= X_train, Y= y_train, m= 50)
θ = pm.Deterministic('θ', pm.math.sigmoid(μ))
y = pm.Bernoulli("y", p= θ, observed= y_train)
with bart_model:
trace = pm.sample()
trace.extend(pm.sample_posterior_predictive(trace))
predtest = pmx.bart.predict(trace, rng, X= X_test)
```

After I use `pmx.bart.predict`

on the out-of-sample data, the recovered values look like:

```
array([-6.93663277, -8.30200875, -7.18379557, ..., -5.97973494,
-6.49198275, -7.58192513])
```

However, after I modify the `pmx.bart.predict`

source code with `return sigmoid(pred)`

, the result is more like I would expect:

```
array([0.00398453, 0.0002981 , 0.00223829, ..., 0.00137673, 0.00213677,
0.00066914])
```

Moving on to my actual application, I have again have a dataset with five (5) variables and a response variable with five classes `array([0, 1, 2, 3, 4, 5]`

. So, I tried to extend the above to a multi-class problem using a softmax function.

```
with pm.Model() as bart_model:
μ = pmx.BART("μ", X= X_train, Y= y_train, m= 50)
θ = pm.math.softmax(μ)
y = pm.Categorical("y", p= θ, observed= y_train)
with bart_model:
trace = pm.sample()
trace.extend(pm.sample_posterior_predictive(trace))
```

Unfortunately, when I check the posterior_predictive, it doesn’t return the ‘y’ class values. Instead, It returns an array of single y values that are obviously not the classes or softmax output.

```
trace.posterior_predictive.y.values
```

```
array([[[3841, 2703, 4008, ..., 2289, 426, 3936],
[6172, 1810, 7215, ..., 1460, 3219, 4301],
[3542, 3660, 2105, ..., 5984, 230, 5654],
...,
[7081, 2198, 5241, ..., 5605, 7490, 4982],
[7121, 6629, 4655, ..., 2991, 7219, 1199],
[ 975, 4058, 3640, ..., 5188, 6736, 6479]],
...,
```

And the subsequent prediction on hold-out data:

```
with bart_model:
predtest = pmx.bart.predict(trace, rng, X= X_test)
```

Returns an array with results that do not appear to have any class structure.

```
array([[[-58.23271276],
[-58.79261905],
[-56.67507199],
...,
[-57.04217796],
[-62.01896868],
[-56.81891283]]])
```

Any help here would be greatly appreciated! Thanks in advance.