type unstable code in Gibbs fails with Enzyme

[penelopeysm]

using Turing, ADTypes, Enzyme, Random, AbstractMCMC

adtype = AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse), function_annotation=Enzyme.Const)
@model function f()
    a ~ Normal()
    b ~ Normal(a)
end

model = f()
spl = Gibbs(
    @varname(a) => HMC(0.1, 10; adtype=adtype),
    @varname(b) => HMC(0.1, 10; adtype=adtype),
)
AbstractMCMC.step(Xoshiro(468), model, spl)

[penelopeysm]

i don't really know where the type instability is and it looks like it might be in dppl

[penelopeysm]

minimise

using Turing, DynamicPPL, ADTypes, Enzyme, Random, AbstractMCMC, LogDensityProblems

adtype = AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse), function_annotation=Enzyme.Const)
@model function f()
    a ~ Normal()
    b ~ Normal(a)
end
model = f()
v = DynamicPPL.VarInfo(model)
gci = Turing.Inference.GibbsContext([@varname(a)], Ref(v), DynamicPPL.DefaultContext());
conditioned_model = DynamicPPL.setleafcontext(model, gci);
vnew = DynamicPPL.VarInfo(conditioned_model);
vnew = DynamicPPL.link(vnew, conditioned_model);
ldf = DynamicPPL.LogDensityFunction(conditioned_model, DynamicPPL.getlogjoint_internal, vnew; adtype=adtype);
LogDensityProblems.logdensity_and_gradient(ldf, randn(1))

[penelopeysm]

The issue is in this if/else block:

Turing.jl/src/mcmc/gibbs.jl

Lines 177 to 206 in bfa61d8

	return if is_target_varname(context, vn)
	# Fall back to the default behavior.
	DynamicPPL.tilde_assume!!(child_context, right, vn, vi)
	elseif has_conditioned_gibbs(context, vn)
	# This branch means that a different sampler is supposed to handle this
	# variable. From the perspective of this sampler, this variable is
	# conditioned on, so we can just treat it as an observation.
	# The only catch is that the value that we need is to be obtained from
	# the global VarInfo (since the local VarInfo has no knowledge of it).
	# Note that tilde_observe!! will trigger resampling in particle methods
	# for variables that are handled by other Gibbs component samplers.
	val = get_conditioned_gibbs(context, vn)
	DynamicPPL.tilde_observe!!(child_context, right, val, vn, vi)
	else
	# If the varname has not been conditioned on, nor is it a target variable, its
	# presumably a new variable that should be sampled from its prior. We need to add
	# this new variable to the global `varinfo` of the context, but not to the local one
	# being used by the current sampler.
	value, new_global_vi = DynamicPPL.tilde_assume!!(
	# child_context might be a PrefixContext so we have to be careful to not
	# overwrite it.
	DynamicPPL.setleafcontext(child_context, DynamicPPL.InitContext()),
	right,
	vn,
	get_global_varinfo(context),
	)
	set_global_varinfo!(context, new_global_vi)
	value, vi
	end
	end

With the example above, vn can either be @varname(a) or @varname(b), so the third branch is never hit. However, for some reason this isn't detected and Enzyme thinks that it can enter the third branch, which has InitContext inside it. Inside InitContext there is a call to push! which is what makes Enzyme complain.

Changing the InitContext to DefaultContext makes the error go away (not saying it's the solution, but just proof that this is where the problem is).

I am not sure why this happens, because @code_typed is_target_varname(gci, @varname(a)) is a compile-time constant::

julia> @code_typed Turing.Inference.is_target_varname(gci, @varname(a))
CodeInfo(
1 ─     return true
) => Bool

julia> @code_typed Turing.Inference.is_target_varname(gci, @varname(b))
CodeInfo(
1 ─     return false
) => Bool

[penelopeysm]

But actually all the code is type stable, so this looks more like an Enzyme problem.

@code_warntype DynamicPPL.tilde_assume!!(gci, Normal(), @varname(a), vnew)

@code_warntype DynamicPPL.tilde_assume!!(InitContext(), Normal(), @varname(a), gci.global_varinfo[])

[penelopeysm]

to be continued another time

[wsmoses]

if you use the code_typed from the error enzyme generates, does that shed any light?

[penelopeysm]

@wsmoses That was the first thing I tried, unfortunately it didn't say much (I didn't look closely at every line but nothing highlighted as red).

using Turing, DynamicPPL, ADTypes, Enzyme, Random, AbstractMCMC, LogDensityProblems

adtype = AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse), function_annotation=Enzyme.Const);
@model function f()
    a ~ Normal()
    b ~ Normal(a)
end;
model = f();
v = DynamicPPL.VarInfo(model);
gci = Turing.Inference.GibbsContext((@varname(a),), Ref(v), DynamicPPL.DefaultContext());
conditioned_model = DynamicPPL.setleafcontext(model, gci);
vnew = DynamicPPL.VarInfo(conditioned_model);
vnew = DynamicPPL.link(vnew, conditioned_model);
ldf = DynamicPPL.LogDensityFunction(conditioned_model, DynamicPPL.getlogjoint_internal, vnew; adtype=adtype);

try
    LogDensityProblems.logdensity_and_gradient(ldf, randn(1))
catch e
    code_typed(e)
end

# -------------------------------


1-element Vector{Any}:
 CodeInfo(
1 ── %1   = $(Expr(:foreigncall, :(:jl_alloc_genericmemory), Ref{Memory{Float64}}, svec(Any, Int64), 0, :(:ccall), Memory{Float64}, 1, 1))::Memory{Float64}
│    %2   = Core.memoryrefnew(%1)::MemoryRef{Float64}
│    %3   = %new(Vector{Float64}, %2, (1,))::Vector{Float64}
│    %4   = Base.getfield(%3, :ref)::MemoryRef{Float64}
│    %5   = Base.memoryrefnew(%4, 1, false)::MemoryRef{Float64}
│           Base.memoryrefset!(%5, r, :not_atomic, false)::Float64
│    %7   = Base.getfield(meta, :idcs)::Dict{VarName{:b, typeof(identity)}, Int64}
│    %8   = Base.getfield(%7, :count)::Int64
│    %9   = Base.add_int(%8, 1)::Int64
│    %10  = Base.getfield(meta, :idcs)::Dict{VarName{:b, typeof(identity)}, Int64}
│           invoke Base.setindex!(%10::Dict{VarName{:b, typeof(identity)}, Int64}, %9::Int64, vn::VarName{:b, typeof(identity)})::Any
│    %12  = Base.getfield(meta, :vns)::Vector{VarName{:b, typeof(identity)}}
│    %13  = Base.getfield(%12, :ref)::MemoryRef{VarName{:b, typeof(identity)}}
│    %14  = Base.getfield(%13, :mem)::Memory{VarName{:b, typeof(identity)}}
│    %15  = Base.getfield(%14, :length)::Int64
│    %16  = Base.getfield(%12, :size)::Tuple{Int64}
│    %17  = $(Expr(:boundscheck, true))::Bool
│    %18  = Base.getfield(%16, 1, %17)::Int64
│    %19  = Base.add_int(%18, 1)::Int64
│    %20  = Base.memoryrefoffset(%13)::Int64
│    %21  = Core.tuple(%19)::Tuple{Int64}
│           Base.setfield!(%12, :size, %21)::Tuple{Int64}
│    %23  = Base.add_int(%20, %19)::Int64
│    %24  = Base.sub_int(%23, 1)::Int64
│    %25  = Base.slt_int(%15, %24)::Bool
└───        goto #3 if not %25
2 ── %27  = %new(Base.var"#133#134"{Vector{VarName{:b, typeof(identity)}}, Int64, Int64, Int64, Int64, Int64, Memory{VarName{:b, typeof(identity)}}, MemoryRef{VarName{:b, typeof(identity)}}}, %12, %24, %20, %19, %18, %15, %14, %13)::Base.var"#133#134"{Vector{VarName{:b, typeof(identity)}}, Int64, Int64, Int64, Int64, Int64, Memory{VarName{:b, typeof(identity)}}, MemoryRef{VarName{:b, typeof(identity)}}}
└───        invoke %27()::MemoryRef{VarName{:b, typeof(identity)}}
3 ┄─        goto #4
4 ── %30  = Base.getfield(%12, :size)::Tuple{Int64}
│    %31  = $(Expr(:boundscheck, true))::Bool
│    %32  = Base.getfield(%30, 1, %31)::Int64
│    %33  = Base.getfield(%12, :ref)::MemoryRef{VarName{:b, typeof(identity)}}
│    %34  = Base.memoryrefnew(%33, %32, false)::MemoryRef{VarName{:b, typeof(identity)}}
│           Base.memoryrefset!(%34, vn, :not_atomic, false)::VarName{:b, typeof(identity)}
└───        goto #5
5 ── %37  = Base.getfield(meta, :vals)::Vector{Float64}
│    %38  = Base.getfield(%37, :size)::Tuple{Int64}
│    %39  = $(Expr(:boundscheck, true))::Bool
│    %40  = Base.getfield(%38, 1, %39)::Int64
│    %41  = Base.getfield(%3, :size)::Tuple{Int64}
│    %42  = $(Expr(:boundscheck, true))::Bool
│    %43  = Base.getfield(%41, 1, %42)::Int64
│    %44  = Base.getfield(meta, :ranges)::Vector{UnitRange{Int64}}
│    %45  = Base.add_int(%40, 1)::Int64
│    %46  = Base.add_int(%40, %43)::Int64
│    %47  = Base.sle_int(%45, %46)::Bool
└───        goto #7 if not %47
6 ──        goto #8
7 ── %50  = Base.sub_int(%45, 1)::Int64
└───        goto #8
8 ┄─ %52  = φ (#6 => %46, #7 => %50)::Int64
│    %53  = %new(UnitRange{Int64}, %45, %52)::UnitRange{Int64}
└───        goto #9
9 ──        goto #10
10 ─ %56  = Base.getfield(%44, :ref)::MemoryRef{UnitRange{Int64}}
│    %57  = Base.getfield(%56, :mem)::Memory{UnitRange{Int64}}
│    %58  = Base.getfield(%57, :length)::Int64
│    %59  = Base.getfield(%44, :size)::Tuple{Int64}
│    %60  = $(Expr(:boundscheck, true))::Bool
│    %61  = Base.getfield(%59, 1, %60)::Int64
│    %62  = Base.add_int(%61, 1)::Int64
│    %63  = Base.memoryrefoffset(%56)::Int64
│    %64  = Core.tuple(%62)::Tuple{Int64}
│           Base.setfield!(%44, :size, %64)::Tuple{Int64}
│    %66  = Base.add_int(%63, %62)::Int64
│    %67  = Base.sub_int(%66, 1)::Int64
│    %68  = Base.slt_int(%58, %67)::Bool
└───        goto #12 if not %68
11 ─ %70  = %new(Base.var"#133#134"{Vector{UnitRange{Int64}}, Int64, Int64, Int64, Int64, Int64, Memory{UnitRange{Int64}}, MemoryRef{UnitRange{Int64}}}, %44, %67, %63, %62, %61, %58, %57, %56)::Base.var"#133#134"{Vector{UnitRange{Int64}}, Int64, Int64, Int64, Int64, Int64, Memory{UnitRange{Int64}}, MemoryRef{UnitRange{Int64}}}
└───        invoke %70()::MemoryRef{UnitRange{Int64}}
12 ┄        goto #13
13 ─ %73  = Base.getfield(%44, :size)::Tuple{Int64}
│    %74  = $(Expr(:boundscheck, true))::Bool
│    %75  = Base.getfield(%73, 1, %74)::Int64
│    %76  = Base.getfield(%44, :ref)::MemoryRef{UnitRange{Int64}}
│    %77  = Base.memoryrefnew(%76, %75, false)::MemoryRef{UnitRange{Int64}}
│           Base.memoryrefset!(%77, %53, :not_atomic, false)::UnitRange{Int64}
└───        goto #14
14 ─ %80  = Base.getfield(meta, :vals)::Vector{Float64}
│    %81  = Base.getfield(%3, :size)::Tuple{Int64}
│    %82  = $(Expr(:boundscheck, true))::Bool
│    %83  = Base.getfield(%81, 1, %82)::Int64
│    %84  = Base.sle_int(0, %83)::Bool
└───        goto #18 if not %84
15 ─ %86  = Base.getfield(%80, :ref)::MemoryRef{Float64}
│    %87  = Base.getfield(%86, :mem)::Memory{Float64}
│    %88  = Base.getfield(%87, :length)::Int64
│    %89  = Base.getfield(%80, :size)::Tuple{Int64}
│    %90  = $(Expr(:boundscheck, true))::Bool
│    %91  = Base.getfield(%89, 1, %90)::Int64
│    %92  = Base.add_int(%91, %83)::Int64
│    %93  = Base.memoryrefoffset(%86)::Int64
│    %94  = Core.tuple(%92)::Tuple{Int64}
│           Base.setfield!(%80, :size, %94)::Tuple{Int64}
│    %96  = Base.add_int(%93, %92)::Int64
│    %97  = Base.sub_int(%96, 1)::Int64
│    %98  = Base.slt_int(%88, %97)::Bool
└───        goto #17 if not %98
16 ─ %100 = %new(Base.var"#133#134"{Vector{Float64}, Int64, Int64, Int64, Int64, Int64, Memory{Float64}, MemoryRef{Float64}}, %80, %97, %93, %92, %91, %88, %87, %86)::Base.var"#133#134"{Vector{Float64}, Int64, Int64, Int64, Int64, Int64, Memory{Float64}, MemoryRef{Float64}}
└───        invoke %100()::MemoryRef{Float64}
17 ┄        goto #19
18 ─ %103 = invoke Base.ArgumentError("grow requires delta >= 0"::String)::ArgumentError
│           Base.throw(%103)::Union{}
└───        unreachable
19 ─ %106 = Base.getfield(%80, :size)::Tuple{Int64}
│    %107 = $(Expr(:boundscheck, true))::Bool
│    %108 = Base.getfield(%106, 1, %107)::Int64
│    %109 = Base.sub_int(%108, %83)::Int64
│    %110 = Base.add_int(%109, 1)::Int64
│    %111 = (%83 === 0)::Bool
└───        goto #21 if not %111
20 ─        goto #53
21 ─ %114 = Base.slt_int(0, %83)::Bool
└───        goto #52 if not %114
22 ─ %116 = $(Expr(:boundscheck, true))::Bool
└───        goto #32 if not %116
23 ─ %118 = Base.add_int(%110, %83)::Int64
│    %119 = Base.sub_int(%118, 1)::Int64
│    %120 = Base.sle_int(%110, %119)::Bool
└───        goto #25 if not %120
24 ─        goto #26
25 ─ %123 = Base.sub_int(%110, 1)::Int64
└───        goto #26
26 ┄ %125 = φ (#24 => %119, #25 => %123)::Int64
│    %126 = %new(UnitRange{Int64}, %110, %125)::UnitRange{Int64}
└───        goto #27
27 ─        goto #28
28 ─ %129 = Core.tuple(%126)::Tuple{UnitRange{Int64}}
│    %130 = Base.getfield(%80, :size)::Tuple{Int64}
│    %131 = $(Expr(:boundscheck, true))::Bool
│    %132 = Base.getfield(%130, 1, %131)::Int64
│    %133 = Base.slt_int(%125, %110)::Bool
│    %134 = Base.sub_int(%110, 1)::Int64
│    %135 = Base.bitcast(UInt64, %134)::UInt64
│    %136 = Base.bitcast(UInt64, %132)::UInt64
│    %137 = Base.ult_int(%135, %136)::Bool
│    %138 = Base.sub_int(%125, 1)::Int64
│    %139 = Base.bitcast(UInt64, %138)::UInt64
│    %140 = Base.bitcast(UInt64, %132)::UInt64
│    %141 = Base.ult_int(%139, %140)::Bool
│    %142 = Base.and_int(%137, %141)::Bool
│    %143 = Base.or_int(%133, %142)::Bool
└───        goto #30 if not %143
29 ─        goto #31
30 ─        invoke Base.throw_boundserror(%80::Vector{Float64}, %129::Tuple{UnitRange{Int64}})::Union{}
└───        unreachable
31 ─        nothing::Nothing
32 ┄ %149 = $(Expr(:boundscheck, true))::Bool
└───        goto #42 if not %149
33 ─ %151 = Base.add_int(1, %83)::Int64
│    %152 = Base.sub_int(%151, 1)::Int64
│    %153 = Base.sle_int(1, %152)::Bool
└───        goto #35 if not %153
34 ─        goto #36
35 ─        goto #36
36 ┄ %157 = φ (#34 => %152, #35 => 0)::Int64
│    %158 = %new(UnitRange{Int64}, 1, %157)::UnitRange{Int64}
└───        goto #37
37 ─        goto #38
38 ─ %161 = Core.tuple(%158)::Tuple{UnitRange{Int64}}
│    %162 = Base.getfield(%3, :size)::Tuple{Int64}
│    %163 = $(Expr(:boundscheck, true))::Bool
│    %164 = Base.getfield(%162, 1, %163)::Int64
│    %165 = Base.slt_int(%157, 1)::Bool
│    %166 = Base.bitcast(UInt64, %164)::UInt64
│    %167 = Base.ult_int(0x0000000000000000, %166)::Bool
│    %168 = Base.sub_int(%157, 1)::Int64
│    %169 = Base.bitcast(UInt64, %168)::UInt64
│    %170 = Base.bitcast(UInt64, %164)::UInt64
│    %171 = Base.ult_int(%169, %170)::Bool
│    %172 = Base.and_int(%167, %171)::Bool
│    %173 = Base.or_int(%165, %172)::Bool
└───        goto #40 if not %173
39 ─        goto #41
40 ─        invoke Base.throw_boundserror(%3::Vector{Float64}, %161::Tuple{UnitRange{Int64}})::Union{}
└───        unreachable
41 ─        nothing::Nothing
42 ┄ %179 = Base.getfield(%80, :ref)::MemoryRef{Float64}
│    %180 = $(Expr(:boundscheck, false))::Bool
│    %181 = Core.memoryrefnew(%179, %110, %180)::MemoryRef{Float64}
│    %182 = Base.getfield(%3, :ref)::MemoryRef{Float64}
│    %183 = $(Expr(:boundscheck, false))::Bool
│    %184 = Core.memoryrefnew(%182, 1, %183)::MemoryRef{Float64}
│    %185 = $(Expr(:gc_preserve_begin, :(%181)))
│    %186 = $(Expr(:gc_preserve_begin, :(%184)))
│    %187 = Base.getfield(%181, :ptr_or_offset)::Ptr{Nothing}
│    %188 = Base.getfield(%184, :ptr_or_offset)::Ptr{Nothing}
│    %189 = Base.mul_int(%83, 8)::Int64
│    %190 = Core.lshr_int(%189, 63)::Int64
│    %191 = Core.trunc_int(Core.UInt8, %190)::UInt8
│    %192 = Core.eq_int(%191, 0x01)::Bool
└───        goto #44 if not %192
43 ─        invoke Core.throw_inexacterror(:convert::Symbol, UInt64::Vararg{Any}, %189)::Union{}
└───        unreachable
44 ─        goto #45
45 ─ %197 = Core.bitcast(Core.UInt64, %189)::UInt64
└───        goto #46
46 ─        goto #47
47 ─        goto #48
48 ─        goto #49
49 ─        $(Expr(:foreigncall, :(:memmove), Ptr{Nothing}, svec(Ptr{Nothing}, Ptr{Nothing}, UInt64), 0, :(:ccall), :(%187), :(%188), :(%197), :(%197), :(%188), :(%187)))::Ptr{Nothing}
└───        goto #50
50 ─        $(Expr(:gc_preserve_end, :(%186)))
│           $(Expr(:gc_preserve_end, :(%185)))
└───        goto #51
51 ─        goto #53
52 ─        invoke Base._throw_argerror("Number of elements to copy must be non-negative."::String)::Union{}
└───        unreachable
53 ┄        goto #54
54 ─        goto #55
55 ─ %212 = Base.getfield(meta, :dists)::Vector{Normal{Float64}}
│    %213 = Base.getfield(%212, :ref)::MemoryRef{Normal{Float64}}
│    %214 = Base.getfield(%213, :mem)::Memory{Normal{Float64}}
│    %215 = Base.getfield(%214, :length)::Int64
│    %216 = Base.getfield(%212, :size)::Tuple{Int64}
│    %217 = $(Expr(:boundscheck, true))::Bool
│    %218 = Base.getfield(%216, 1, %217)::Int64
│    %219 = Base.add_int(%218, 1)::Int64
│    %220 = Base.memoryrefoffset(%213)::Int64
│    %221 = Core.tuple(%219)::Tuple{Int64}
│           Base.setfield!(%212, :size, %221)::Tuple{Int64}
│    %223 = Base.add_int(%220, %219)::Int64
│    %224 = Base.sub_int(%223, 1)::Int64
│    %225 = Base.slt_int(%215, %224)::Bool
└───        goto #57 if not %225
56 ─ %227 = %new(Base.var"#133#134"{Vector{Normal{Float64}}, Int64, Int64, Int64, Int64, Int64, Memory{Normal{Float64}}, MemoryRef{Normal{Float64}}}, %212, %224, %220, %219, %218, %215, %214, %213)::Base.var"#133#134"{Vector{Normal{Float64}}, Int64, Int64, Int64, Int64, Int64, Memory{Normal{Float64}}, MemoryRef{Normal{Float64}}}
└───        invoke %227()::MemoryRef{Normal{Float64}}
57 ┄        goto #58
58 ─ %230 = Base.getfield(%212, :size)::Tuple{Int64}
│    %231 = $(Expr(:boundscheck, true))::Bool
│    %232 = Base.getfield(%230, 1, %231)::Int64
│    %233 = Base.getfield(%212, :ref)::MemoryRef{Normal{Float64}}
│    %234 = Base.memoryrefnew(%233, %232, false)::MemoryRef{Normal{Float64}}
│           Base.memoryrefset!(%234, dist, :not_atomic, false)::Normal{Float64}
└───        goto #59
59 ─ %237 = Base.getfield(meta, :is_transformed)::BitVector
│    %238 = Base.getfield(%237, :chunks)::Vector{UInt64}
│    %239 = Base.getfield(%237, :len)::Int64
│    %240 = Base.and_int(%239, 63)::Int64
│    %241 = (%240 === 0)::Bool
└───        goto #69 if not %241
60 ─ %243 = Base.getfield(%238, :ref)::MemoryRef{UInt64}
│    %244 = Base.getfield(%243, :mem)::Memory{UInt64}
│    %245 = Base.getfield(%244, :length)::Int64
│    %246 = Base.getfield(%238, :size)::Tuple{Int64}
│    %247 = $(Expr(:boundscheck, true))::Bool
│    %248 = Base.getfield(%246, 1, %247)::Int64
│    %249 = Base.add_int(%248, 1)::Int64
│    %250 = Base.memoryrefoffset(%243)::Int64
│    %251 = Core.tuple(%249)::Tuple{Int64}
│           Base.setfield!(%238, :size, %251)::Tuple{Int64}
│    %253 = Base.add_int(%250, %249)::Int64
│    %254 = Base.sub_int(%253, 1)::Int64
│    %255 = Base.slt_int(%245, %254)::Bool
└───        goto #62 if not %255
61 ─ %257 = %new(Base.var"#133#134"{Vector{UInt64}, Int64, Int64, Int64, Int64, Int64, Memory{UInt64}, MemoryRef{UInt64}}, %238, %254, %250, %249, %248, %245, %244, %243)::Base.var"#133#134"{Vector{UInt64}, Int64, Int64, Int64, Int64, Int64, Memory{UInt64}, MemoryRef{UInt64}}
└───        invoke %257()::MemoryRef{UInt64}
62 ┄        goto #63
63 ─ %260 = Base.getfield(%238, :size)::Tuple{Int64}
│    %261 = $(Expr(:boundscheck, true))::Bool
│    %262 = Base.getfield(%260, 1, %261)::Int64
│    %263 = $(Expr(:boundscheck, true))::Bool
└───        goto #67 if not %263
64 ─ %265 = Base.sub_int(%262, 1)::Int64
│    %266 = Base.bitcast(UInt64, %265)::UInt64
│    %267 = Base.getfield(%238, :size)::Tuple{Int64}
│    %268 = $(Expr(:boundscheck, true))::Bool
│    %269 = Base.getfield(%267, 1, %268)::Int64
│    %270 = Base.bitcast(UInt64, %269)::UInt64
│    %271 = Base.ult_int(%266, %270)::Bool
└───        goto #66 if not %271
65 ─        goto #67
66 ─ %274 = Core.tuple(%262)::Tuple{Int64}
│           invoke Base.throw_boundserror(%238::Vector{UInt64}, %274::Tuple{Int64})::Union{}
└───        unreachable
67 ┄ %277 = Base.getfield(%238, :ref)::MemoryRef{UInt64}
│    %278 = Base.memoryrefnew(%277, %262, false)::MemoryRef{UInt64}
│           Base.memoryrefset!(%278, 0x0000000000000000, :not_atomic, false)::UInt64
└───        goto #68
68 ─        nothing::Nothing
69 ┄ %282 = Base.getfield(%237, :len)::Int64
│    %283 = Base.add_int(%282, 1)::Int64
│           Base.setfield!(%237, :len, %283)::Int64
└───        goto #70
70 ─        return DynamicPPL.nothing
) => Nothing

[penelopeysm]

using Turing, DynamicPPL, Enzyme

@model function f()
    a ~ Normal()
    b ~ Normal(a)
end
model = f()
v = DynamicPPL.VarInfo(model)
vref = Ref(v)
function f(x)
    vi = last(DynamicPPL.tilde_assume!!(InitContext(), Normal(x[1]), @varname(b), vref[]))
    return DynamicPPL.getlogjoint(vi)
end

f(randn(1))

Enzyme.gradient(set_runtime_activity(Enzyme.Reverse), f, randn(1))

[penelopeysm]

using DynamicPPL, Enzyme

@model function f()
    a ~ Normal()
    b ~ Normal(a)
end
model = f()
v = VarInfo(model)
vref = Ref(v)

function tasm(vn, dist, vi)
    x = rand(dist)
    return if haskey(vi, vn)
        DynamicPPL.setindex!!(vi, x, vn)
    else
        DynamicPPL.push!!(vi, vn, x, dist)
    end
end
function f(x)
    vi = vref[]
    vn = @varname(b)
    dist = Normal(x[1])
    vi = tasm(vn, dist, vi)
    return DynamicPPL.getlogjoint(vi)
end

f(randn(1))

Enzyme.gradient(set_runtime_activity(Enzyme.Reverse), f, randn(1))

[penelopeysm]

using DynamicPPL, Distributions, Enzyme

@model function f()
    a ~ Normal()
    b ~ Normal(a)
end
model = f()
v = VarInfo(model)
vref = Ref(v)

function tasm(vn, vi)
    # this function must not be inlined
    return if haskey(vi, vn)
        vi
    else
        push!(vi.metadata.b.vns, vn)
        vi
    end
end
function f(x)
    vi = vref[]
    vn = @varname(b)
    vi = tasm(vn, vi)
    return DynamicPPL.getlogjoint(vi)
end

f(randn(1))

Enzyme.gradient(set_runtime_activity(Enzyme.Reverse), f, randn(1))

[penelopeysm]

using AbstractPPL, Enzyme

struct MD{Tvn<:VarName}
    vns::Vector{Tvn}
end
struct VI{Tmd}
    metadata::Tmd
    lp::Float64
end
function Base.haskey(vi, vn)
    return any(md -> vn in md.vns, values(vi.metadata))
end
v = VI((; b = MD([@varname(b)])), 0.0)
vref = Ref(v)
@assert haskey(v, @varname(b))

function tasm(vn, vi)
    # this function must not be inlined
    return if haskey(vi, vn)
        vi
    else
        push!(vi.metadata.b.vns, vn)
        vi
    end
end
function f(x)
    vi = vref[]
    vn = @varname(b)
    vi = tasm(vn, vi)
    return vi.lp
end

f(randn(1))

Enzyme.gradient(set_runtime_activity(Enzyme.Reverse), f, randn(1))

[penelopeysm]

Minimised and filed upstream EnzymeAD/Enzyme.jl#2738