control_flow model runs unreliably

[penelopeysm]

The purpose of the control_flow model is to demonstrate errors with compiled ReverseDiff:

using DynamicPPL, LogDensityProblems, Distributions, ADTypes, ForwardDiff, ReverseDiff

@model function control_flow()
    a ~ Normal()
    if a > 0
        b ~ Normal()
    else
        b ~ Beta(2, 2)
    end
end

# Gradient preparation
model = control_flow()
adtype = AutoReverseDiff(; compile=true)
vi = DynamicPPL.unflatten(VarInfo(model), [0.5, -0.5]) # note: a > 0 and b < 0
ldf = LogDensityFunction(model, vi; adtype=adtype)

# For comparison
ldf_ref = LogDensityFunction(model, vi; adtype=AutoForwardDiff())

# Two sets of parameters
p1 = [0.5, -0.5]  # a > 0, b < 0
p2 = [-0.5, 0.5]  # a < 0, b > 0

julia> LogDensityProblems.logdensity_and_gradient(ldf, p1)
(-2.0878770664093453, [-0.5, 0.5])

julia> LogDensityProblems.logdensity_and_gradient(ldf_ref, p1)
(-2.0878770664093453, [-0.5, 0.5])

julia> LogDensityProblems.logdensity_and_gradient(ldf, p2)
(-2.0878770664093453, [0.5, -0.5])

julia> LogDensityProblems.logdensity_and_gradient(ldf_ref, p2)
(-0.6384734250965083, [0.5, 0.0])

As can be seen from the above, the gradient evaluates correctly at p1, but incorrectly at p2. The issue is that reproducing this is contingent on choosing the right set parameters to pass to the model; if we pass p1 it looks like compiled ReverseDiff runs fine.

[penelopeysm]

Funnily enough, this is where the idea of a DITest scenario is perfect because the scenario also encapsulates x

[penelopeysm]

Although (as far as I can tell) it's not currently possible to control the values used to prepare the gradient in a Scenario JuliaDiff/DifferentiationInterface.jl#771

[penelopeysm]

Fixed here by 1874554 and 7ac92d2