Thank you, it is working now with a slight modification of adding θ to posterior_mean.sel({"θ_dim_1":i}).θ.values
posterior_mean = az.extract_dataset(trace, var_names='θ').mean("sample")
new_y = [np.random.multinomial(n=1, pvals=posterior_mean.sel({"θ_dim_1":i}).θ.values).argmax() for i in range(7500)]
I was interested in using predicting out-of-sample data, but using pmx.bart.predict doesn’t seem to be doing a good job with unbalanced classes.