Enzyme reverse stochastic failures with `Threads.@threads` + TSVI

[penelopeysm]

Run this extract using #1113 or any subsequent PR. The correct gradient is -51.0 (both FiniteDifferences and ForwardDiff get this right). However, sometimes Enzyme misses out on some terms and reports -50 or -48. This happens with both slow LDF and fast LDF, so it's unrelated to #1113.

(Make sure to launch Julia with multiple threads, of course)

using DynamicPPL, Distributions, LogDensityProblems, Chairmarks, LinearAlgebra, ADTypes, ForwardDiff, Enzyme, FiniteDifferences

if Threads.nthreads() == 1
    error("run this code with multiple threads")
end

const adtypes = (
    AutoFiniteDifferences(; fdm=central_fdm(5, 1)),
    AutoForwardDiff(),
    AutoEnzyme(; mode=set_runtime_activity(Reverse), function_annotation=Const)
)

function check_ldfs(model)
    vi = VarInfo(model)
    xs = [1.0]
    for adtype in adtypes
        sldf = DynamicPPL.LogDensityFunction(model, getlogjoint, vi; adtype=adtype)
        fldf = DynamicPPL.Experimental.FastLDF(model, getlogjoint, vi; adtype=adtype)
        sldf_grad = LogDensityProblems.logdensity_and_gradient(sldf, xs)
        fldf_grad = LogDensityProblems.logdensity_and_gradient(fldf, xs)
        @show adtype
        @show sldf_grad
        @show fldf_grad
        @assert sldf_grad[2] ≈ fldf_grad[2]
    end
end
@model function threads(y=zeros(50))
    x ~ Normal()
    Threads.@threads for i in eachindex(y)
        y[i] ~ Normal(x)
    end
end
check_ldfs(setthreadsafe(threads(), true))

[penelopeysm]

It also refuses to differentiate these:

@model function threads2(y=zeros(50))
    x ~ Normal()
    loglikes = map(y) do yi
        Threads.@spawn begin
            return logpdf(Normal(x), yi)
        end
    end
    @addlogprob! sum(fetch.(loglikes))
end
check_ldfs(threads2())

prints

 Warning: active variables passed by value to jl_new_task are not yet supported
└ @ Enzyme.Compiler ~/.julia/packages/Enzyme/EP5KJ/src/rules/parallelrules.jl:726
┌ Warning: active variables passed by value to jl_new_task are not yet supported
└ @ Enzyme.Compiler ~/.julia/packages/Enzyme/EP5KJ/src/rules/parallelrules.jl:726

before it actually errors.

Same with:

using OhMyThreads
@model function threads3(y=zeros(50))
    x ~ Normal()
    let x = x
        lls = tmapreduce(+, eachindex(y)) do i
            logpdf(Normal(x), y[i])
        end
        @addlogprob! lls
    end
end
check_ldfs(threads3())

[yebai]

This reminds me of the days when we struggled and eventually gave up on Zygote.

[penelopeysm]

DI's docs say:

https://juliadiff.org/DifferentiationInterface.jl/DifferentiationInterface/stable/api/#DifferentiationInterface.prepare_gradient

The preparation result prep is not thread-safe. Sharing it between threads may lead to unexpected behavior. If you need to run differentiation concurrently, prepare separate prep objects for each thread.

So, technically, the way LogDensityFunction is constructed now is not correct either

[mhauru]

Eek, and that's true for all backends. Another reason to make TSVI opt-in, and make any use of threading more of a "know what you're doing, there be dragons here" thing?

[gdalle]

So, technically, the way LogDensityFunction is constructed now is not correct either

Wait, can you clarify what you mean?
If you have multithreading inside the function you differentiate, you're fine. What you're not allowed to do is something like this, which differentiates the same function from different threads while mutating the same prep object:

prep = prepare_gradient(f, backend, x)
@threads for i in 1:10
    gradient(f, prep, backend, xs[i])
end

[gdalle]

JuliaDiff/DifferentiationInterface.jl#923

[penelopeysm]

Reported upstream EnzymeAD/Enzyme.jl#3005

[penelopeysm]

Update: fixed in Enzyme v0.13.140