`dppl_sto_volatility` correctness issue with multiple backends

[penelopeysm]

using DynamicPPL
using DynamicPPL.TestUtils.AD: run_ad
using Distributions
using ADTypes: AutoForwardDiff, AutoMooncake
import ForwardDiff
import Mooncake
using Random: Xoshiro

y = [1502.39, 1505.62, 1509.48, 1507.72, 1512.58]

@model function dppl_sto_volatility(y, ::Type{Tv}=Vector{Float64}) where {Tv}
    T = length(y)
    μ ~ Cauchy(0, 10)
    ϕ ~ Uniform(-1, 1)
    σ ~ truncated(Cauchy(0, 5); lower=0)
    h = Tv(undef, T)
    h[1] ~ Normal(μ, σ / sqrt(1 - ϕ^2))
    y[1] ~ Normal(0, exp(h[1] / 2))
    for t in 2:T
        h[t] ~ Normal(μ + ϕ * (h[t-1] - μ), σ)
        y[t] ~ Normal(0, exp(h[t] / 2))
    end
end

model = dppl_sto_volatility(y[1:5])
xs = DynamicPPL.link(VarInfo(Xoshiro(468), model), model)[:]

fd_res = run_ad(model, AutoForwardDiff(); test=false, params=xs)

mc_res = run_ad(model, AutoMooncake(); test=false, params=xs)

fd_res.grad_actual .- mc_res.grad_actual
#=
8-element Vector{Float64}:
  0.0
  2.7755575615628914e-17
  5.551115123125783e-16
 -1.9342813113834067e25
  0.0
  1.9342813113834067e25
  0.0
 -1.9342813113834067e25
=#

[penelopeysm]

here is a smaller repro

using DynamicPPL
using DynamicPPL.TestUtils.AD: run_ad
using Distributions
using ADTypes: AutoForwardDiff, AutoMooncake
import ForwardDiff
import Mooncake
using Random: Xoshiro

@model function f(y)
    m ~ Cauchy(0, 10)
    s ~ truncated(Cauchy(0, 5); lower=0)
    h ~ Normal(m, s)
    y ~ Normal(0, exp(h / 2))
end
model = f(1505.62)
xs = DynamicPPL.link(VarInfo(Xoshiro(468), model), model)[:]

fd_res = run_ad(model, AutoForwardDiff(); test=false, params=xs).grad_actual

mc_res = run_ad(model, AutoMooncake(); test=false, params=xs).grad_actual

fd_res .- mc_res

it's worth noting that actually the huge difference is a tiny relative difference because the gradient is on the order of 1e41:

julia> fd_res = run_ad(model, AutoForwardDiff(); test=false, params=xs).grad_actual
[ Info: Running AD on f with AutoForwardDiff()
       params : [-80.49419773440104, -1.7776148807021088, -80.65985398063759]
       actual : (-1.2148712356518068e41, [-5.772813893107857, 0.9580727253507426, 1.2148712356518068e41])
3-element Vector{Float64}:
 -5.772813893107857
  0.9580727253507426
  1.2148712356518068e41

julia> mc_res = run_ad(model, AutoMooncake(); test=false, params=xs).grad_actual
[ Info: Running AD on f with AutoMooncake{Nothing}(nothing)
       params : [-80.49419773440104, -1.7776148807021088, -80.65985398063759]
       actual : (-1.2148712356518068e41, [-5.772813893107856, 0.9580727253507425, 1.2148712356518066e41])
3-element Vector{Float64}:
 -5.772813893107856
  0.9580727253507425
  1.2148712356518066e41

julia> fd_res .- mc_res
3-element Vector{Float64}:
 -8.881784197001252e-16
  1.1102230246251565e-16
  1.9342813113834067e25

which suggests that maybe the way we test for equality (currently using atol) is not correct and we should be testing rtol on a per-element basis.

[penelopeysm]

Fixed by TuringLang/DynamicPPL.jl#964 but we will have to wait for DPPL 0.37 to be released before we can use it.