[Bug]: `infer()` fails when a prior is passed as distribution and AddonMemory() is activated

[skoghoern]

Description

When defining a @model and passing a prior as a distribution type, the infer() function fails when including addons = (AddonMemory(),). infer() works fine without AddonMemory().
When defining the prior inside the @model it works correctly with AddonMemory().

Expected Behavior

infer() should run including AddonMemory when priors passed as distribution.

Actual Behavior

Failed with error (see below).
It seems like there is no Addon attached to the message sent from the prior.

Minimal Reproducible Example or steps to reproduce the issue

@model function outside(prior_node,obs)
    p ~ prior_node
    obs ~ Categorical(p)
end
@model function inside(prior,obs)
    p ~ Dirichlet([prior, 1.0])
    obs ~ Categorical(p)
end
result = infer(
        model=outside(prior_node=Dirichlet([1.0, 1.0])), # fails with AddonMemory
        # model=inside(prior=1.0), # works fine with AddonMemory
        data=(obs=[1.0,0.0],),
        addons = (AddonMemory(),)
)

Error Message / Stack Trace

MethodError: no method matching multiply_addons(::Nothing, ::Tuple{AddonMemory{ReactiveMP.AddonMemoryMessageMapping{ReactiveMP.MessageMapping{Categorical{P} where P<:Real, Val{:p}, Marginalisation, Nothing, Val{(:out,)}, Nothing, Tuple{AddonMemory{Nothing}}, Nothing, Nothing}, Nothing, Tuple{PointMass{Vector{Float64}}}, Dirichlet{Float64, Vector{Float64}, Float64}}}}, ::Dirichlet{Float64, Vector{Float64}, Float64}, ::Dirichlet{Float64, Vector{Float64}, Float64}, ::Dirichlet{Float64, Vector{Float64}, Float64})
The function `multiply_addons` exists, but no method is defined for this combination of argument types.

Closest candidates are:
  multiply_addons(::Nothing, ::Nothing, ::Any, ::Any, ::Any)
   @ ReactiveMP C:\Users\BMW\.julia\packages\ReactiveMP\N8Y9O\src\addons.jl:7
  multiply_addons(::Nothing, ::Any, ::Any, ::Missing, ::Any)
   @ ReactiveMP C:\Users\BMW\.julia\packages\ReactiveMP\N8Y9O\src\addons.jl:8
  multiply_addons(::Tuple, ::Tuple, ::Any, ::Any, ::Any)
   @ ReactiveMP C:\Users\BMW\.julia\packages\ReactiveMP\N8Y9O\src\addons.jl:14
  ...


Stacktrace:
  [1] multiply_messages(prod_strategy::GenericProd, left::Message{Dirichlet{Float64, Vector{Float64}, Float64}, Nothing}, right::Message{Dirichlet{Float64, Vector{Float64}, Float64}, Tuple{AddonMemory{ReactiveMP.AddonMemoryMessageMapping{ReactiveMP.MessageMapping{Categorical{P} where P<:Real, Val{:p}, Marginalisation, Nothing, Val{(:out,)}, Nothing, Tuple{AddonMemory{Nothing}}, Nothing, Nothing}, Nothing, Tuple{PointMass{Vector{Float64}}}, Dirichlet{Float64, Vector{Float64}, Float64}}}}})
    @ ReactiveMP C:\Users\BMW\.julia\packages\ReactiveMP\N8Y9O\src\message.jl:134
  [2] (::ReactiveMP.var"#16#18"{GenericProd})(left::Message{Dirichlet{Float64, Vector{Float64}, Float64}, Nothing}, right::Message{Dirichlet{Float64, Vector{Float64}, Float64}, Tuple{AddonMemory{ReactiveMP.AddonMemoryMessageMapping{ReactiveMP.MessageMapping{Categorical{P} where P<:Real, Val{:p}, Marginalisation, Nothing, Val{(:out,)}, Nothing, Tuple{AddonMemory{Nothing}}, Nothing, Nothing}, Nothing, Tuple{PointMass{Vector{Float64}}}, Dirichlet{Float64, Vector{Float64}, Float64}}}}})
    @ ReactiveMP C:\Users\BMW\.julia\packages\ReactiveMP\N8Y9O\src\message.jl:150

Julia Version

1.11

RxInfer Version

Latest stable

Environment Information

- Package Versions: [86711068] RxInfer v4.7.0

Session ID (Optional)

No response

Additional Context

No response

[bvdmitri]

Confirmed, thanks for the report!

[gauraaansh]

Quick update on this:

The exact reproduction above doesn't apply anymore since AddonMemory was removed in the v6 annotations rewrite (it just throws a deprecation error now pointing at InputArgumentsAnnotations). But I tried the same thing with the new API and it still breaks:

@model function outside(prior_node, obs)
    p ~ prior_node
    obs ~ Categorical(p)
end

infer(
    model = outside(prior_node = Dirichlet([1.0, 1.0])),
    data = (obs = [1.0, 0.0],),
    options = (annotations = (InputArgumentsAnnotations(),),),
)

KeyError: key :rule_input_arguments not found, coming from post_product_annotations! in input_arguments.jl. Same pattern as before: prior defined inline works, prior passed in as a Distribution doesn't.

Turns out the cause is: a clamped/constant message (like a prior passed in directly) never goes through a rule, so its annotation dict stays empty. post_product_annotations! assumes both sides of a product always have one, so it blows up when one side is clamped. The existing tests for that function always populate both sides manually, so this case was never actually tested.

I've got a fix for it (just checking which side has the annotation before merging instead of assuming both do), and I ran the whole ReactiveMP.jl suite against it (still all passing), added tests for the missing-annotation cases, and confirmed the actual posteriors don't change, just the bookkeeping.

Also found the same issue independently in LogScaleAnnotations, but that seems unrelated/pre-existing so I left it out of scope here.

Let me know if this sounds right and I'll open a PR.

cc @bvdmitri, since you looked at this one originally

[bvdmitri]

a PR would be very much welcome! @gauraaansh