Enzyme doesn't like accumulators

[penelopeysm]

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