Warning: TypeAnalysisDepthLimit, incorrect gradient

[tpapp]

The uninferred version of #840 runs, but produces an incorrect gradient. Enzyme and julia versions same as in #840.

output

julia> _, y = Enzyme.autodiff(Enzyme.ReverseWithPrimal, ℓ, Enzyme.Active, Enzyme.Duplicated(x, ∂ℓ_∂x))

┌ Warning: TypeAnalysisDepthLimit
│   store {} addrspace(10)* %3, {} addrspace(10)* addrspace(10)* %7, align 8, !dbg !19, !tbaa !36, !alias.scope !26, !noalias !29
│ {[]:Pointer, [0]:Pointer, [0,0]:Pointer, [0,0,0]:Pointer, [0,0,0,0]:Pointer, [0,0,0,0,0]:Pointer, [0,0,0,0,0,0]:Pointer, [0,0,0,1,0]:Pointer, [0,0,0,1,0,0]:Pointer, [0,0,0,2,0]:Pointer, [0,0,0,2,0,0]:Pointer, [0,0,0,3,0]:Pointer, [0,0,0,3,0,0]:Pointer, [0,0,0,4,0]:Pointer, [0,0,0,4,0,0]:Pointer, [0,0,0,5,0]:Pointer, [0,0,0,5,0,0]:Pointer, [0,0,0,6,0]:Pointer, [0,0,0,6,0,0]:Pointer, [0,0,0,7,0]:Pointer, [0,0,0,7,0,0]:Pointer, [0,1,0]:Pointer, [0,1,0,0]:Pointer, [0,2,0]:Pointer, [0,2,0,0]:Pointer, [0,3,0]:Pointer, [0,3,0,0]:Pointer, [0,4,0]:Pointer, [0,4,0,0]:Pointer, [0,5,0]:Pointer, [0,5,0,0]:Pointer, [0,6,0]:Pointer, [0,6,0,0]:Pointer, [0,7,0]:Pointer, [0,7,0,0]:Pointer}
│ 
│ Stacktrace:
│  [1] #log_posterior#5
│    @ ./REPL[6]:2
└ @ Enzyme.Compiler ~/.julia/packages/GPUCompiler/BxfIW/src/utils.jl:56
┌ Warning: TypeAnalysisDepthLimit
│   store atomic {} addrspace(10)* %6, {} addrspace(10)* addrspace(11)* %5 release, align 8, !dbg !19, !tbaa !20, !alias.scope !26, !noalias !29
│ {[]:Pointer, [0]:Pointer, [0,0]:Pointer, [0,0,0]:Pointer, [0,0,0,0]:Pointer, [0,0,0,0,0]:Pointer, [0,0,0,0,0,0]:Pointer, [0,0,0,0,1,0]:Pointer, [0,0,0,0,2,0]:Pointer, [0,0,0,0,3,0]:Pointer, [0,0,0,0,4,0]:Pointer, [0,0,0,0,5,0]:Pointer, [0,0,0,0,6,0]:Pointer, [0,0,0,0,7,0]:Pointer, [0,0,1,0]:Pointer, [0,0,1,0,0]:Pointer, [0,0,2,0]:Pointer, [0,0,2,0,0]:Pointer, [0,0,3,0]:Pointer, [0,0,3,0,0]:Pointer, [0,0,4,0]:Pointer, [0,0,4,0,0]:Pointer, [0,0,5,0]:Pointer, [0,0,5,0,0]:Pointer, [0,0,6,0]:Pointer, [0,0,6,0,0]:Pointer, [0,0,7,0]:Pointer, [0,0,7,0,0]:Pointer}
│ 
│ Stacktrace:
│  [1] #log_posterior#5
│    @ ./REPL[6]:2
└ @ Enzyme.Compiler ~/.julia/packages/GPUCompiler/BxfIW/src/utils.jl:56

MWE

using StaticArrays, ForwardDiff
import Enzyme

const ω_DIMENSION = 7

function log_posterior(; ω_intercept, ω_std)
    _uni_prior(z::Real) = -log(5) - abs2(z / 5)/2
    _uni_prior(z::AbstractVector) = sum(_uni_prior, z)
    ℓ_prior = _uni_prior(ω_intercept) + _uni_prior(ω_std)
    ℓ_prior
end

function transform_and_logjac(x)
    S = eltype(x)
    logjac::S = 0.0
    i::Int = 0
    function _pick(d::Int)
        v = @view x[(i+1):(i+d)]
        i += d
        v
    end
    ω_intercept = SVector{ω_DIMENSION}(_pick(ω_DIMENSION))
    pre_ω_std = SVector{ω_DIMENSION}(_pick(ω_DIMENSION))
    ω_std = exp.(pre_ω_std)
    logjac += sum(pre_ω_std)
    (;  ω_intercept, ω_std), logjac
end

function ℓ(x)
    θ, logjac = transform_and_logjac(x)
    log_posterior(; θ...) + logjac
end # NOTE this function is NOT INFERRED

const D = 2 * ω_DIMENSION
x = zeros(D)

Enzyme.API.runtimeActivity!(true)
∂ℓ_∂x = zero(x)
_, y = Enzyme.autodiff(Enzyme.ReverseWithPrimal, ℓ, Enzyme.Active, Enzyme.Duplicated(x, ∂ℓ_∂x))

# comparison

∂ℓ_∂x_FD = ForwardDiff.gradient(ℓ, x)
err = abs.(∂ℓ_∂x_FD .- ∂ℓ_∂x)
err_ix = findall(err .≥ 1e-6)
@show err_ix
@show err[err_ix]

[wsmoses]

MWE:

using StaticArrays
using Enzyme

Enzyme.API.printall!(true)

@noinline function log_posterior(w_intercept)
	first(w_intercept)
end

function func(x)
    t = SVector{1}(Float64[x])
    log_posterior(Base.inferencebarrier(t))::Float64
end

y = Enzyme.autodiff(Enzyme.Reverse, func, Active(2.0))
@show y