Type error in backward pass

[Red-Portal]

Hi, here's the MWE for the failing test in AdvancedVI.

using Bijectors
using DifferentiationInterface
using Distributions
using Mooncake
using Optimisers
using Random

function f(params, aux)
    (; re, samples, baseline, q_stop) = aux
    q = re(params)

    ℓq = logpdf.(Ref(q), eachcol(samples))
    ℓq_stop = logpdf.(Ref(q_stop), eachcol(samples))
    ℓπ = sum(abs2, samples, dims=1)[1,:]
    ℓπ_mean = mean(ℓπ)
    score_grad = mean(@. ℓq * (ℓπ - baseline))
    score_grad_stop = mean(@. ℓq_stop * (ℓπ - baseline))
    energy = ℓπ_mean + (score_grad - score_grad_stop)
    energy
end

function main()
    rng = Random.default_rng()
    q0  = MvNormal(zeros(3), ones(3))
    b   = Bijectors.Stacked(
        Bijectors.bijector.([LogNormal(0, 1), MvNormal(zeros(2), ones(2))]),
        [1:1, 2:3]
    )
    q  = Bijectors.transformed(q0, Bijectors.inverse(b), )

    params, re = Optimisers.destructure(q)
    
    adtype = AutoMooncake(; config=nothing)
    aux = (
        samples  = rand(rng, q, 10),
        baseline = 1.0,
        re       = re,
        q_stop   = q,
    )
    value_and_gradient(f, adtype, params, Constant(aux))
end

[willtebbutt]

Thanks for narrowing this down. I won't have time to look at it today unforunately, but I should do on Monday.

[willtebbutt]

Alas, the emergence of v1.11 in CI has meant that I need to focus on that. I'll try to ensure that this is resolved by the end of the week.

[Red-Portal]

The v0.3 release of AdvancedVI will wait until this is resolved!

[willtebbutt]

Annoyingly the 1.11 transition is taking longer than I had planned, so this isn't going to get a look until next week. Apologies!

[Red-Portal]

I saw all the works that you're going through. No apology needed! And thanks for all the effort.

[willtebbutt]

@Red-Portal I believe that this is resolved in 0.4.18. Please let me know if it is, and close this issue if you find that it is!

[willtebbutt]

edit: I meant 0.4.17

[Red-Portal]

Running the tests just now! Thanks for all the work.

[Red-Portal]

@willtebbutt The original error seems to have disappeared, but a new one took its place. Until I can package a MWE, here is the error.

[willtebbutt]

Thanks for this. I'll take a look later today!

[willtebbutt]

@Red-Portal this problem should be fixed on 0.4.20 . If you could please let me know, and close this issue if it is resolved.

[yebai]

This now works on Julia 1.11, but it is still broken for Julia 1.0, see https://github.com/TuringLang/AdvancedVI.jl/actions/runs/11459971691/job/31885647008?pr=99#step:6:831

[willtebbutt]

@yebai this should work on 0.4.24, which will be available shortly. Could you re-start CI once it's released, and let me know if the problem has been resolved?

[yebai]

Thanks @willtebbutt -- it now works.

The remaining errors in TuringLang/AdvancedVI.jl#99 are no longer associated with Mooncake.