Multi-class BART Model Assistance

Questions v5
4

Thank you - I tried the below, but it did not work.

with pm.Model(coords=COORDS) as bart_model:
    μ = pmx.BART("μ", X= X_train, Y= y_train, m= 50, shape=(6, 7500))
    #θ = pm.Deterministic('θ', pm.math.softmax(μ))
    θ = pm.math.softmax(μ)
    y = pm.Categorical("y", p= θ, observed= y_train)

Which threw a really long error summarized with:

AssertionError: Could not broadcast dimensions