Enzyme throws lots of errors when using MLDatasets

[penelopeysm]

That's why, for example, https://turinglang.org/ADTests/pr/ (built from #23) currently shows a lot of errors.

Currently minimised to:

module MWE

using DynamicPPL: DynamicPPL, VarInfo, @model
import Enzyme: Enzyme, set_runtime_activity, Forward, Reverse, Const
using Distributions: InverseGamma, MvNormal, product_distribution
using LinearAlgebra: Diagonal

@model function demo_assume_multivariate_observe(x = [1.5, 2.0])
    s ~ product_distribution([InverseGamma(2, 3), InverseGamma(2, 3)])
    x ~ MvNormal([0.0, 0.0], Diagonal(s))
end

# Removing any of these three blocks causes it to run fine
struct PStruct
    a
end
Base.getproperty(::Type{PStruct}, s::Symbol) = getfield(PStruct, s)
struct QStruct
    a
end
Base.getproperty(::Type{QStruct}, s::Symbol) = getfield(QStruct, s)
struct RStruct
    a
end
Base.getproperty(::Type{RStruct}, s::Symbol) = getfield(RStruct, s)
# =========================================

model = demo_assume_multivariate_observe()
varinfo = VarInfo(model)
params = varinfo[:]

function f(
    x::AbstractVector, model::DynamicPPL.Model, varinfo::DynamicPPL.AbstractVarInfo
)
    varinfo_new = DynamicPPL.unflatten(varinfo, x)
    _, vi = DynamicPPL.evaluate!!(model, varinfo_new, DynamicPPL.DefaultContext())
    return vi.logp[]
end

Enzyme.gradient(set_runtime_activity(Forward), f, params, Const(model), Const(varinfo))

end

[penelopeysm]

module MWE

using ChangesOfVariables: ChangesOfVariables
import Enzyme: Enzyme, Reverse, Const, set_runtime_activity

Enzyme.Compiler.VERBOSE_ERRORS[] = true

struct PStruct
    p
end
Base.getproperty(::Type{PStruct}, s::Symbol) = getfield(PStruct, s)
struct QStruct
    q
end
Base.getproperty(::Type{QStruct}, s::Symbol) = getfield(QStruct, s)
struct RStruct
    r
end
Base.getproperty(::Type{RStruct}, s::Symbol) = getfield(RStruct, s)

# The below minimised from:
# @model function f()
#     s ~ product_distribution([InverseGamma(2, 3), InverseGamma(2, 3)])
#     1.5 ~ Normal(s[1])
# end

function wladj(mapped_f::Base.Fix1{typeof(broadcast)}, X)
    f = mapped_f.x
    y_with_ladj = broadcast(Base.Fix1(ChangesOfVariables.with_logabsdet_jacobian, f), X)
    ChangesOfVariables._with_ladj_on_mapped(broadcast, y_with_ladj)
end

function f(x::AbstractVector)
    logp = Ref(0.0)
    g = Base.Fix1(broadcast, Base.Fix1(broadcast, exp))
    s, _ = wladj(g, x)
    logp[] += sum(s)
    return logp[]
end

x = [0.5, 1.0]
@show f(x)

Enzyme.gradient(Reverse, f, x)

end

[penelopeysm]

module MWE

import Enzyme: Enzyme, Reverse, Const, set_runtime_activity

Enzyme.Compiler.VERBOSE_ERRORS[] = true

struct PStruct
    p
end
Base.getproperty(::Type{PStruct}, s::Symbol) = getfield(PStruct, s)
struct QStruct
    q
end
Base.getproperty(::Type{QStruct}, s::Symbol) = getfield(QStruct, s)
struct RStruct
    r
end
Base.getproperty(::Type{RStruct}, s::Symbol) = getfield(RStruct, s)

# The below minimised from:
# @model function f()
#     s ~ product_distribution([InverseGamma(2, 3), InverseGamma(2, 3)])
#     1.5 ~ Normal(s[1])
# end

function g(xs)
    return broadcast(x -> (exp(x), x), xs)
end

function f(x::AbstractVector)
    logp = Ref(0.0)
    ts = broadcast(g, x)
    logp[] += (ts[1][1] + ts[2][1])
    return logp[]
end

x = [0.5, 1.0]
@show f(x)

Enzyme.gradient(Reverse, f, x)

end

[penelopeysm]

Upstreamed EnzymeAD/Enzyme.jl#2408