Enzyme doesn't like accumulators

[penelopeysm]

Index of upstreamed issues

✅ EnzymeAD/Enzyme.jl#2429
✅ EnzymeAD/Enzyme.jl#2512
✅ EnzymeAD/Enzyme.jl#2518
✅ EnzymeAD/Enzyme.jl#2580
✅ (fixed on DPPL side) EnzymeAD/Enzyme.jl#2589

---

Original comment

MWE (make sure to use breaking branch, or something with accumulators)

using DynamicPPL, Distributions
using ADTypes: AutoEnzyme
import Enzyme: Enzyme, Forward, Reverse, set_runtime_activity, Const

@model f() = x ~ Normal()
model = f()
vi = VarInfo(model)
ctx = DefaultContext()
params = [0.5]

function logp(x, model, vi, ctx)
    varinfo_new = DynamicPPL.unflatten(vi, x)
    varinfo_eval = last(DynamicPPL.evaluate!!(model, varinfo_new, ctx))
    return getlogjoint(varinfo_eval)
end

logp(params, model, vi, ctx)

Enzyme.gradient(set_runtime_activity(Forward), logp, params, Const(model), Const(vi), Const(ctx))

Enzyme.gradient(set_runtime_activity(Reverse), logp, params, Const(model), Const(vi), Const(ctx))

[penelopeysm]

Upstreamed EnzymeAD/Enzyme.jl#2429

[penelopeysm]

Sort of related to #944 -- ultimately the problematic part was the call to deepcopy(getaccs(vi)) inside unflatten.

If I replace deepcopy(getaccs(vi)) with just vi.accs, Enzyme runs fine, and CI also passes: #948 --- although I'm not sure if this is a general solution because it doesn't account for potentially having non-isbitstype accumulators.

DynamicPPL.jl/src/varinfo.jl

Lines 437 to 455 in 80db9e2

	function unflatten(vi::VarInfo, x::AbstractVector)
	md = unflatten_metadata(vi.metadata, x)
	# Use of float_type_with_fallback(eltype(x)) is necessary to deal with cases where x is
	# a gradient type of some AD backend.
	# TODO(mhauru) How could we do this more cleanly? The problem case is map_accumulator!!
	# for ThreadSafeVarInfo. In that one, if the map produces e.g a ForwardDiff.Dual, but
	# the accumulators in the VarInfo are plain floats, we error since we can't change the
	# element type of ThreadSafeVarInfo.accs_by_thread. However, doing this conversion here
	# messes with cases like using Float32 of logprobs and Float64 for x. Also, this is just
	# plain ugly and hacky.
	# The below line is finicky for type stability. For instance, assigning the eltype to
	# convert to into an intermediate variable makes this unstable (constant propagation)
	# fails. Take care when editing.
	accs = map(
	acc -> convert_eltype(float_type_with_fallback(eltype(x)), acc),
	deepcopy(getaccs(vi)),
	)
	return VarInfo(md, accs)
	end

[penelopeysm]

The immediate bug is fixed, but that just reveals a ton more issues.

[penelopeysm]

Current progress. Reverse fails, forward is fine.

I am starting to think that removing ThreadSafeVarInfo would be a much better use of my time than wading through the accumulator code for it, which has been a really frustrating experience.

The example below contains no DynamicPPL code except for the accumulators themselves and AccumulatorTuple constructor.

using DynamicPPL, Distributions
import Enzyme: Enzyme, Forward, Reverse, set_runtime_activity, Const
@model f() = x ~ Normal()
model = f()
vi = VarInfo(model)
vi = DynamicPPL.setaccs!!(vi, (DynamicPPL.LogPriorAccumulator(0.0), DynamicPPL.LogLikelihoodAccumulator(0.0)))
params = [0.5]

# struct LP{T<:Real} <: DynamicPPL.LogProbAccumulator{T}
#     logp::T
# end
# DynamicPPL.logp(l::LP) = l.logp
# DynamicPPL.accumulator_name(::Type{<:LP}) = :LogPrior
# function DynamicPPL.accumulate_assume!!(acc::LP, val, logjac, vn, right)
#     return DynamicPPL.acclogp(acc, logpdf(right, val))
# end
# DynamicPPL.accumulate_observe!!(acc::LP, right, left, vn) = acc
# struct LL{T<:Real} <: DynamicPPL.LogProbAccumulator{T}
#     logp::T
# end
# DynamicPPL.logp(l::LL) = l.logp
# DynamicPPL.accumulator_name(::Type{<:LL}) = :LogLikelihood
# DynamicPPL.accumulate_assume!!(acc::LL, val, logjac, vn, right) = acc
# function DynamicPPL.accumulate_observe!!(acc::LL, right, left, vn)
#     return DynamicPPL.acclogp(acc, Distributions.loglikelihood(right, left))
# end
# needs at least two accs for it to fail
# vi = DynamicPPL.setaccs!!(vi, (LP(0.0), LL(0.0)))

m(::DynamicPPL.LogPriorAccumulator) = :LogPrior
nm(::DynamicPPL.LogLikelihoodAccumulator) = :LogLikelihood
struct AT{NT}
    nt::NT
end
function toat(nt)
    t = values(nt)
    names = map(nm, t)
    return AT(NamedTuple{names}(t))
end
struct V3{Tat}
    accs::Tat
end
vi = V3(DynamicPPL.AccumulatorTuple((; LogPrior = DynamicPPL.LogPriorAccumulator(0.0), LogLikelihood = DynamicPPL.LogLikelihoodAccumulator(0.0))))
struct TSVI4{V,L<:AT}
    varinfo::V
    accs_by_thread::Vector{L}
end
function TSVI4(vi)
    return TSVI4(vi, [AT(vi.accs.nt)])
end
function gacc(vi::TSVI4, ::Val{accname_inner}) where {accname_inner}
    main_acc = vi.varinfo.accs.nt[accname_inner]
    other_accs = map(accs -> accs.nt[accname_inner], vi.accs_by_thread)
    return foldl(DynamicPPL.combine, other_accs; init=main_acc)
end
@noinline id(x) = x
# function acas(acc::LP, x)
#     LP(acc.logp + id(x))
# end
# acas(acc::LL, x) = acc
function acas(acc::DynamicPPL.LogPriorAccumulator, x)
    DynamicPPL.LogPriorAccumulator(acc.logp + id(x))
end
acas(acc::DynamicPPL.LogLikelihoodAccumulator, x) = acc
function mdlf(x, vi)
    f = acc -> acas(acc, x)
    new_nt = map(f, vi.accs_by_thread[1].nt)
    vi.accs_by_thread[1] = AT(new_nt)
    return vi
end
function gacs(tsvi)
    accname_vals = map(Val, keys(tsvi.varinfo.accs))
    return toat(map(anv -> gacc(tsvi, anv), accname_vals)).nt
end
function evlt(x, vi)
    return if Threads.nthreads() > 1
        tsvi = TSVI4(vi)
        tsvi = mdlf(x, id(tsvi))
        gacs(tsvi)
    else
        nt = vi.accs.nt
        new_val = DynamicPPL.LogPriorAccumulator(nt[:LogPrior].logp + x)
        merge(nt, (; LogPrior = new_val))
    end
end
function f(x, vi)
    nt = evlt(x[1], vi)
    return nt.LogPrior.logp
end

f(params, vi)

Enzyme.gradient(set_runtime_activity(Reverse), f, params, Const(vi))

Enzyme.gradient(set_runtime_activity(Forward), f, params, Const(vi))

[penelopeysm]

no DynamicPPL

import Enzyme: Enzyme, Forward, Reverse, set_runtime_activity, Const

abstract type AbsT{T<:Real} end
basetypeof(::T) where {T} = Base.typename(T).wrapper
function cmbn(acc::AbsT, acc2::AbsT)
    return basetypeof(acc)(acc.logp + acc2.logp)
end

struct LP{T<:Real} <: AbsT{T}
    logp::T
end
struct LL{T<:Real} <: AbsT{T}
    logp::T
end
nm(::LP) = :LogPrior
nm(::LL) = :LogLikelihood

struct AT2{N,T<:NamedTuple}
    nt::T
end
function AT2(t::T) where {N,T<:NTuple{N,Any}}
    names = map(nm, t)
    nt = NamedTuple{names}(t)
    return AT2{N,typeof(nt)}(nt)
end
AT2(nt::NamedTuple) = AT2(tuple(nt...))
Base.keys(at::AT2) = keys(at.nt)

struct V0{Tat}
    accs::Tat
end
# needs at least two accs for it to fail
vi = V0(AT2((; LogPrior = LP(0.0), LogLikelihood = LL(0.0))))
struct TSVI4{V,L<:AT2}
    varinfo::V
    accs_by_thread::Vector{L}
end
function TSVI4(vi)
    return TSVI4(vi, [AT2(vi.accs.nt)])
end
function gacc(vi::TSVI4, ::Val{accname_inner}) where {accname_inner}
    main_acc = vi.varinfo.accs.nt[accname_inner]
    other_accs = map(accs -> accs.nt[accname_inner], vi.accs_by_thread)
    return foldl(cmbn, other_accs; init=main_acc)
end
@noinline id(x) = x
acas(acc::LP, x) = LP(acc.logp + x)
acas(acc::LL, x) = acc
function mdlf(x, vi)
    f = acc -> acas(acc, x)
    new_nt = map(f, vi.accs_by_thread[1].nt)
    vi.accs_by_thread[1] = AT2(new_nt)
    return vi
end
function gacs(tsvi)
    accname_vals = map(Val, keys(tsvi.varinfo.accs))
    return AT2(map(anv -> gacc(tsvi, anv), accname_vals)).nt
end
function evlt(x, vi)
    return if Threads.nthreads() > 1
        tsvi = TSVI4(vi)
        tsvi = mdlf(x, id(tsvi))
        gacs(tsvi)
    else
        nt = vi.accs.nt
        new_val = LP(nt[:LogPrior].logp + x)
        merge(nt, (; LogPrior = new_val))
    end
end
function f(x, vi)
    nt = evlt(x[1], vi)
    return nt.LogPrior.logp
end

params = [0.5]
f(params, vi)
Enzyme.gradient(set_runtime_activity(Reverse), f, params, Const(vi))

[penelopeysm]

EnzymeAD/Enzyme.jl#2512

[penelopeysm]

Completely different error (again forward-mode is fine):

using DynamicPPL, Distributions
import Enzyme: Enzyme, Forward, Reverse, set_runtime_activity, Const

@model function f(x=[1.5, 2.0])
    s = Vector{Float64}(undef, length(x))
    m = Vector{Float64}(undef, length(x))
    for i in eachindex(s)
        s[i] ~ InverseGamma(2, 3)
    end
    m ~ Normal()
    x[1] ~ Normal(m, s[1])
end
model = f()
vi = VarInfo(model)
params = fill(0.5, 4)
function logp(x, model, vi)
    varinfo_new = DynamicPPL.unflatten(vi, x)
    varinfo_eval = last(DynamicPPL.evaluate!!(model, varinfo_new))
    return getlogjoint(varinfo_eval)
end
logp(params, model, vi)
Enzyme.gradient(set_runtime_activity(Reverse), logp, params, Const(model), Const(vi))

[penelopeysm]

Exciting times because this just crashes outright instead of erroring. All the DynamicPPL stuff is gone now, the only 'special' thing left is VarName.

using Accessors: IndexLens, PropertyLens
using AbstractPPL: VarName
import Enzyme: Enzyme, Forward, Reverse, set_runtime_activity, Const

#= This is still not enough to replicate the crash
const ALLOWED_OPTICS = Union{typeof(identity),PropertyLens,IndexLens,ComposedFunction}
struct VN{sym,T<:ALLOWED_OPTICS}
    inner::T
    VN{sym}(inner) where {sym} = new{sym,typeof(inner)}(inner)
end
=#

struct VI{Tmd,Tlogp}
    metadata::Tmd
    logp::Tlogp
end
struct TSVI{V}
    varinfo::V
end
struct MyMetadata{
    TIdcs<:Dict{<:VarName,Int},
    TVal<:AbstractVector{<:Real},
}
    idcs::TIdcs
    vals::TVal
end
@noinline function assm(vi)
    y = vi.varinfo.metadata.s.vals[1:1]
    return only(y)
end
function g(vi, z)
    s = Vector{Float64}(undef, length(z))
    for i in eachindex(s)
        s[i] = assm(vi)
    end
    m = vi.varinfo.metadata.m.vals[1:2]
    return VI(vi.varinfo.metadata, 1.0)
end
params = zeros(4)
function f(x)
    md_s = MyMetadata(
        Dict{VarName,Int}(VarName{:s}(IndexLens((1,))) => 1),
        x[1:2],
    )
    md_m = MyMetadata(
        Dict{VarName,Int}(VarName{:m}(identity) => 1),
        x[3:4],
    )
    vi = TSVI(VI((; s=md_s, m=md_m), 0.0))
    vi = g(vi, [1.5, 2.0])
    return vi.logp
end
f(params)
Enzyme.gradient(set_runtime_activity(Reverse), f, params)

[penelopeysm]

EnzymeAD/Enzyme.jl#2518

[penelopeysm]

Still errors, I'm going to cry in a corner

using DynamicPPL, Distributions
import Enzyme: Enzyme, Forward, Reverse, set_runtime_activity, Const
@model function f(x=[1.5, 2.0])
    s = Vector{Float64}(undef, length(x))
    m = Vector{Float64}(undef, length(x))
    for i in eachindex(s)
        s[i] ~ InverseGamma(2, 3)
    end
    m ~ Normal()
    x[1] ~ Normal(m, s[1])
end
vi = VarInfo(f())

function tass(right, vn, vi)
    y = DynamicPPL.getindex_internal(vi, vn)
    x = only(y)
    vi = DynamicPPL.accumulate_assume!!(vi, x, 0.0, vn, right)
    return x, vi
end
function tass_m(right, vi)
    x = only(vi.varinfo.metadata.m.vals)
    vi = DynamicPPL.accumulate_assume!!(vi, x, 0.0, @varname(m), right)
    return x, vi
end
function tobs(right, left, vn, vi)
    vi = DynamicPPL.accumulate_observe!!(vi, right, left, vn)
    return left, vi
end
function modelf(varinfo, x=[1.5, 2.0])
    s = Vector{Float64}(undef, length(x))
    m = Vector{Float64}(undef, length(x))
    for i in eachindex(s)
        s[i], varinfo = tass(InverseGamma(2, 3), @varname(s[i]), varinfo)
    end
    m, varinfo = tass_m(Normal(), varinfo)
    x[1], varinfo = tobs(Normal(m, s[1]), x[1], @varname(x[1]), varinfo)
    return varinfo
end
params = fill(0.5, 4)
function logp(x, vi)
    DynamicPPL.AccumulatorTuple(DynamicPPL.LogLikelihoodAccumulator()))
    varinfo_new = DynamicPPL.unflatten(vi, x)
    tsvi = DynamicPPL.ThreadSafeVarInfo(varinfo_new)
    new_tsvi = modelf(tsvi)
    return getloglikelihood(new_tsvi)
end
logp(params, vi)
Enzyme.gradient(set_runtime_activity(Reverse), logp, params, Const(vi))