Mis-specifying dimensions in params() leads to hard trace when used within RxInfer

Short-story: When using ExponentialFamilyProjection to project a logpdf onto MvNormal using GaussNewton, if we specify the D dimension wrong, we get an error from Distributions.jl. The below MWE shows the error, and the associated stacktrace clearly points to the logpdf() call. But if you then are projecting inside RxInfer infer(), inside prod rules or message outputs, you will have a harder time tracing. See further below "RxInfer version".

Blah blah, haha, it took some time, but I'll suggest a check to add to bring awareness to the error early on.

Minimum working example:

    using RxInfer, ExponentialFamilyProjection, ForwardDiff
    function total_logpdf(x) return logpdf( MvNormalMeanCovariance( [1,2], diageye(2)), x) end
    function my_logpdf!(out, x)
        out[1] = total_logpdf(x)
    end
    function my_grad_hess!(out_grad, out_hess, x)
        out_grad .= ForwardDiff.gradient(total_logpdf, x)
        out_hess .= ForwardDiff.hessian(total_logpdf, x)
    end
    params = ProjectionParameters(
        tolerance = 1e-6,
        strategy = ExponentialFamilyProjection.GaussNewton(nsamples = 0),
    )
    inplace_enzyme = ExponentialFamilyProjection.InplaceLogpdfGradHess(my_logpdf!, my_grad_hess!)
    prj_enzyme = ProjectedTo(MvNormalMeanCovariance, 5; parameters = params) # <-- Note the wrong number of dimensions, 5
    projected_Gauss = project_to(prj_enzyme, inplace_enzyme)

"RxInfer Version"

Offending call:

prod_foldl_reduce(prod_constraint, form_constraint, ::FormConstraintCheckLast) =
    (messages) -> constrain_form_as_message(foldl((left, right) -> multiply_messages(prod_constraint, left, right), Base.Generator(as_message, messages)), form_constraint)

Supporting data:

length(collect(Base.Generator(as_message, messages)))
	2
getdata(as_message(messages[1]))
	MvNormalMeanCovariance(
		μ: [0.06167592982615934, 0.23937339493931517, 0.7356798139497275, 0.0, 0.0, 0.0]
		Σ: [1.0e-6 0.0 0.0 0.0 0.0 0.0; 0.0 1.0e-6 0.0 0.0 0.0 0.0; 0.0 0.0 1.0e-6 0.0 0.0 0.0; 0.0 0.0 0.0 1.0e-6 0.0 0.0; 0.0 0.0 0.0 0.0 1.0e-6 0.0; 0.0 0.0 0.0 0.0 0.0 1.0e-6])
getdata(as_message(messages[2]))
	ContinuousMultivariateLogPdf(UnspecifiedDomain())

Then, computation proceeds into multiply_messages():

function multiply_messages(prod_strategy, left::Message, right::Message)
    # We propagate clamped message, in case if both are clamped
    is_prod_clamped = is_clamped(left) && is_clamped(right)
    # We propagate initial message, in case if both are initial or left is initial and right is clameped or vice-versa
    is_prod_initial = !is_prod_clamped && (is_clamped_or_initial(left)) && (is_clamped_or_initial(right))

    # process distributions
    left_dist  = getdata(left)
    right_dist = getdata(right)
    new_dist   = prod(prod_strategy, left_dist, right_dist)

    # process addons
    left_addons  = getaddons(left)
    right_addons = getaddons(right)

    # process addons
    new_addons = multiply_addons(left_addons, right_addons, new_dist, left_dist, right_dist)

    return Message(new_dist, is_prod_clamped, is_prod_initial, new_addons)
end

Then into prod(). It is not 100% clear that this is the prod rule being called, but there are no others defined for this combination of messages.
Furthermore, calling prod(prod_constraint, as_message(messages[1]), as_message(messages[2])) returns the same error so I can be sure the error is happening there:

function ReactiveMP.prod(::GenericProd, left::TYPE_Norm, right::TYPE_ContLogPdf) 
    function total_logpdf(x) return logpdf(left, x) + logpdf(right, x) end
    function my_logpdf!(out, x)
        out[1] = total_logpdf(x)
    end
    function my_grad_hess!(out_grad, out_hess, x)
        out_grad .= ForwardDiff.gradient(total_logpdf, x)
        out_hess .= ForwardDiff.hessian(total_logpdf, x)
    end
    params = ProjectionParameters(
        tolerance = 1e-6,
        strategy = ExponentialFamilyProjection.GaussNewton(nsamples = 0),
    )
    inplace_enzyme = ExponentialFamilyProjection.InplaceLogpdfGradHess(my_logpdf!, my_grad_hess!)
    prj_enzyme = ProjectedTo(MvNormalMeanCovariance, D; parameters = params)
    projected_Gauss = project_to(prj_enzyme, inplace_enzyme)
    return projected_Gauss
end

At this point I can't trace further because the debugger stops at the above "offending call", but searching for the error message points to the following (logical) logpdf function call in Distributions.jl:

"""
    logpdf(d::Distribution{ArrayLikeVariate{N}}, x::AbstractArray{<:Real,N}) where {N}

Evaluate the logarithm of the probability density function of `d` at `x`.

This function checks if the size of `x` is compatible with distribution `d`. This check can
be disabled by using `@inbounds`.

# Implementation

Instead of `logpdf` one should implement `_logpdf(d, x)` which does not have to check the
size of `x`.

See also: [`pdf`](@ref), [`gradlogpdf`](@ref).
"""
@inline function logpdf(
    d::Distribution{ArrayLikeVariate{N}}, x::AbstractArray{<:Real,M}
) where {N,M}
    if M == N
        @boundscheck begin
            size(x) == size(d) ||
                throw(DimensionMismatch("inconsistent array dimensions"))
        end
        return _logpdf(d, x)
    else
        @boundscheck begin
            M > N ||
                throw(DimensionMismatch(
                    "number of dimensions of the variates ($M) must be greater than or equal to the dimension of the distribution ($N)"
                ))
            ntuple(i -> size(x, i), Val(N)) == size(d) ||
                throw(DimensionMismatch("inconsistent array dimensions"))
        end
        return @inbounds map(Base.Fix1(logpdf, d), eachvariate(x, variate_form(typeof(d))))
    end
end

And so one of the two above was triggering the error. Not sure which, but presumably for the MvNormal dist because otherwise the equality check would fail first at size(ContinuousMultivariateLogpdf) which is undefined.