Multinomial logistic regression: `predict_particles` throws `MethodError`

[DilumAluthge]

predict_particles throws a MethodError when I try to use it in the multinomial logistic regression example.

The error is: ERROR: MethodError: no method matching &(::Particles{Bool,1000}, ::Particles{Bool,1000})

Full output:

julia> include("examples/example-multinomial-logistic-regression.jl");
[ Info: Training Machine{SossMLJModel{,…}} @705.
size(predictor_marginal) = (150,)
size(iris, 1) = 150
Evaluating over 4 folds: 100%[=========================] Time: 0:04:29
Evaluating over 4 folds: 100%[=========================] Time: 0:04:14

julia> Xnew = iris[!, feature_columns]
150×4 DataFrame
│ Row │ PetalLength │ PetalWidth │ SepalLength │ SepalWidth │
│     │ Float64     │ Float64    │ Float64     │ Float64    │
├─────┼─────────────┼────────────┼─────────────┼────────────┤
│ 1   │ 1.4         │ 0.2        │ 5.1         │ 3.5        │
⋮
│ 150 │ 5.1         │ 1.8        │ 5.9         │ 3.0        │

julia> particles = predict_particles(mach, Xnew; response = :y)
ERROR: MethodError: no method matching &(::Particles{Bool,1000}, ::Particles{Bool,1000})
Closest candidates are:
  &(::Any, ::Any, ::Any, ::Any...) at operators.jl:538
  &(::VectorizationBase.Static{M}, ::Number) where M at /Users/dilum/.julia/packages/VectorizationBase/kIoqa/src/static.jl:123
  &(::Number, ::VectorizationBase.Static{N}) where N at /Users/dilum/.julia/packages/VectorizationBase/kIoqa/src/static.jl:124
  ...
Stacktrace:
  [1] _broadcast_getindex_evalf
    @ ./broadcast.jl:648 [inlined]
  [2] _broadcast_getindex
    @ ./broadcast.jl:621 [inlined]
  [3] getindex
    @ ./broadcast.jl:575 [inlined]
  [4] copy
    @ ./broadcast.jl:876 [inlined]
  [5] materialize
    @ ./broadcast.jl:837 [inlined]
  [6] _check_probs_01(probs::Vector{Particles{Float64,1000}})
    @ MLJBase ~/.julia/packages/MLJBase/02f6S/src/univariate_finite/types.jl:66
  [7] _broadcast_getindex_evalf
    @ ./broadcast.jl:648 [inlined]
  [8] _broadcast_getindex
    @ ./broadcast.jl:621 [inlined]
  [9] getindex
    @ ./broadcast.jl:575 [inlined]
 [10] copy
    @ ./broadcast.jl:876 [inlined]
 [11] materialize
    @ ./broadcast.jl:837 [inlined]
 [12] UnivariateFinite(::MLJModelInterface.FullInterface, prob_given_class::OrderedCollections.LittleDict{CategoricalValue{String,UInt8},AbstractVector{Particles{Float64,1000}},Vector{CategoricalValue{String,UInt8}},Vector{AbstractVector{Particles{Float64,1000}}}}; kwargs::Base.Iterators.Pairs{Symbol,Any,Tuple{Symbol,Symbol},NamedTuple{(:pool, :ordered),Tuple{CategoricalPool{String,UInt8,CategoricalValue{String,UInt8}},Bool}}})
    @ MLJBase ~/.julia/packages/MLJBase/02f6S/src/univariate_finite/types.jl:127
 [13] _UnivariateFinite(support::CategoricalVector{String,UInt8,String,CategoricalValue{String,UInt8},Union{}}, probs::Matrix{Particles{Float64,1000}}, N::Int64; augment::Bool, kwargs::Base.Iterators.Pairs{Symbol,Any,Tuple{Symbol,Symbol},NamedTuple{(:pool, :ordered),Tuple{CategoricalPool{String,UInt8,CategoricalValue{String,UInt8}},Bool}}})
    @ MLJBase ~/.julia/packages/MLJBase/02f6S/src/univariate_finite/types.jl:244
 [14] #_UnivariateFinite#36
    @ ~/.julia/packages/MLJBase/02f6S/src/univariate_finite/types.jl:286 [inlined]
 [15] #_UnivariateFinite#39
    @ ~/.julia/packages/MLJBase/02f6S/src/univariate_finite/types.jl:295 [inlined]
 [16] UnivariateFinite(::MLJModelInterface.FullInterface, support::Vector{String}, probs::Matrix{Particles{Float64,1000}}; kwargs::Base.Iterators.Pairs{Symbol,CategoricalPool{String,UInt8,CategoricalValue{String,UInt8}},Tuple{Symbol},NamedTuple{(:pool,),Tuple{CategoricalPool{String,UInt8,CategoricalValue{String,UInt8}}}}})
    @ MLJBase ~/.julia/packages/MLJBase/02f6S/src/univariate_finite/types.jl:211
 [17] UnivariateFinite(support::Vector{String}, probs::Matrix{Particles{Float64,1000}}; kwargs::Base.Iterators.Pairs{Symbol,CategoricalPool{String,UInt8,CategoricalValue{String,UInt8}},Tuple{Symbol},NamedTuple{(:pool,),Tuple{CategoricalPool{String,UInt8,CategoricalValue{String,UInt8}}}}})
    @ MLJModelInterface ~/.julia/packages/MLJModelInterface/Zu24c/src/data_utils.jl:431
 [18] #UnivariateFinite#23
    @ ~/.julia/packages/MLJBase/02f6S/src/univariate_finite/types.jl:33 [inlined]
 [19] macro expansion
    @ ~/.julia/packages/GeneralizedGenerated/wp5nX/src/closure_conv.jl:121 [inlined]
 [20] _particles(#unused#::Type{TypeEncoding(Main)}, _m::Soss.Model{NamedTuple{(:X, :pool, :β),T} where T<:Tuple,TypeEncoding(begin
    η = X * β
    μ = NNlib.softmax(η; dims = 2)
    y_dists = UnivariateFinite(pool.levels, μ; pool = pool)
    n = size(X, 1)
    y ~ For((j->begin
                    y_dists[j]
                end), n)
end),TypeEncoding(Main)}, _args::NamedTuple{(:X, :pool, :β),Tuple{Matrix{Float64},CategoricalPool{String,UInt8,CategoricalValue{String,UInt8}},Matrix{Particles{Float64,1000}}}}, _n::Val{1000})
    @ Soss ~/.julia/packages/GeneralizedGenerated/wp5nX/src/closure_conv.jl:121
 [21] particles
    @ ~/.julia/packages/Soss/FgRzp/src/particles.jl:66 [inlined]
 [22] predict_particles(predictor::SossMLJ.SossMLJPredictor{SossMLJModel{UnivariateFinite,Soss.Model{NamedTuple{(:X, :pool),T} where T<:Tuple,TypeEncoding(begin
    k = length(pool.levels)
    p = size(X, 2)
    β ~ Normal(0.0, 1.0) |> iid(p, k)
    η = X * β
    μ = NNlib.softmax(η; dims = 2)
    y_dists = UnivariateFinite(pool.levels, μ; pool = pool)
    n = size(X, 1)
    y ~ For((j->begin
                    y_dists[j]
                end), n)
end),TypeEncoding(Main)},NamedTuple{(:pool,),Tuple{CategoricalPool{String,UInt8,CategoricalValue{String,UInt8}}}},typeof(dynamicHMC),Symbol,typeof(SossMLJ.default_transform)},Vector{NamedTuple{(:β,),Tuple{Matrix{Float64}}}},Soss.Model{NamedTuple{(:X, :pool, :β),T} where T<:Tuple,TypeEncoding(begin
    η = X * β
    μ = NNlib.softmax(η; dims = 2)
    y_dists = UnivariateFinite(pool.levels, μ; pool = pool)
    n = size(X, 1)
    y ~ For((j->begin
                    y_dists[j]
                end), n)
end),TypeEncoding(Main)},NamedTuple{(:X, :pool),Tuple{Matrix{Float64},CategoricalPool{String,UInt8,CategoricalValue{String,UInt8}}}}}, Xnew::DataFrame)
    @ SossMLJ ~/Downloads/SossMLJ.jl/src/particles.jl:12
 [23] predict_particles(sm::SossMLJModel{UnivariateFinite,Soss.Model{NamedTuple{(:X, :pool),T} where T<:Tuple,TypeEncoding(begin
    k = length(pool.levels)
    p = size(X, 2)
    β ~ Normal(0.0, 1.0) |> iid(p, k)
    η = X * β
    μ = NNlib.softmax(η; dims = 2)
    y_dists = UnivariateFinite(pool.levels, μ; pool = pool)
    n = size(X, 1)
    y ~ For((j->begin
                    y_dists[j]
                end), n)
end),TypeEncoding(Main)},NamedTuple{(:pool,),Tuple{CategoricalPool{String,UInt8,CategoricalValue{String,UInt8}}}},typeof(dynamicHMC),Symbol,typeof(SossMLJ.default_transform)}, fitresult::NamedTuple{(:model, :post),Tuple{SossMLJModel{UnivariateFinite,Soss.Model{NamedTuple{(:X, :pool),T} where T<:Tuple,TypeEncoding(begin
    k = length(pool.levels)
    p = size(X, 2)
    β ~ Normal(0.0, 1.0) |> iid(p, k)
    η = X * β
    μ = NNlib.softmax(η; dims = 2)
    y_dists = UnivariateFinite(pool.levels, μ; pool = pool)
    n = size(X, 1)
    y ~ For((j->begin
                    y_dists[j]
                end), n)
end),TypeEncoding(Main)},NamedTuple{(:pool,),Tuple{CategoricalPool{String,UInt8,CategoricalValue{String,UInt8}}}},typeof(dynamicHMC),Symbol,typeof(SossMLJ.default_transform)},Vector{NamedTuple{(:β,),Tuple{Matrix{Float64}}}}}}, Xnew::DataFrame; response::Symbol)
    @ SossMLJ ~/Downloads/SossMLJ.jl/src/particles.jl:20
 [24] predict_particles(mach::Machine{SossMLJModel{UnivariateFinite,Soss.Model{NamedTuple{(:X, :pool),T} where T<:Tuple,TypeEncoding(begin
    k = length(pool.levels)
    p = size(X, 2)
    β ~ Normal(0.0, 1.0) |> iid(p, k)
    η = X * β
    μ = NNlib.softmax(η; dims = 2)
    y_dists = UnivariateFinite(pool.levels, μ; pool = pool)
    n = size(X, 1)
    y ~ For((j->begin
                    y_dists[j]
                end), n)
end),TypeEncoding(Main)},NamedTuple{(:pool,),Tuple{CategoricalPool{String,UInt8,CategoricalValue{String,UInt8}}}},typeof(dynamicHMC),Symbol,typeof(SossMLJ.default_transform)}}, Xraw::DataFrame; kwargs::Base.Iterators.Pairs{Symbol,Symbol,Tuple{Symbol},NamedTuple{(:response,),Tuple{Symbol}}})
    @ SossMLJ ~/Downloads/SossMLJ.jl/src/machine-operations.jl:10
 [25] top-level scope
    @ REPL[4]:1

---

For comparison, predict_particles works correctly in the linear regression example:

julia> include("examples/example-linear-regression.jl");
[ Info: Training Machine{SossMLJModel{,…}} @946.
Evaluating over 6 folds: 100%[=========================] Time: 0:00:04
Evaluating over 6 folds: 100%[=========================] Time: 0:00:04

julia> particles = predict_particles(mach, X; response = :y)
100-element Vector{Particles{Float64,1000}}:
 -0.0407 ± 0.054
 -0.0283 ± 0.053
  ⋮
 -0.118 ± 0.054
 -0.25 ± 0.055

[DilumAluthge]

Maybe instead of response = :y, I should try response = :μ and/or response = :y_dists.

[DilumAluthge]

response = :μ and response = :y_dists both give the same error as above.

This is curious. If I have response = :μ, I shouldn't even need to construct the UnivariateFinite distributions, right? After all, I'm only asking for μ.

[DilumAluthge]

@cscherrer Is there a way to construct the particles but e.g. tell Soss that "I only want up to μ, so skip any variables that depend on μ"?

[cscherrer]

I think the easiest way to do this would be something like

1. Have Soss make a new model with just the variables you want particles for
2. Get the particles
3. Call predict on those in the original model

For (1), check out Soss.below

[DilumAluthge]

I'm going to mark this as potentially breaking, since it will probably require some changes to the return types of public functions.