AssertionError: RT <: Array

[mhauru]

MWE:

using Enzyme: Enzyme
using LinearAlgebra: Diagonal, transpose
using Turing: Turing

Turing.@model function demo(x=transpose([1.5 2.0;]))
    m ~ Turing.MvNormal(Diagonal([1.0, 1.0]))
    x .~ Turing.MvNormal(m, Diagonal([1.0, 1.0]))
    return nothing
end

model = demo()
ctx = Turing.DefaultContext()
vi = Turing.VarInfo(model)
ldp = Turing.LogDensityFunction(vi, model, ctx)
f = Base.Fix1(Turing.LogDensityProblems.logdensity, ldp)
x = [0.0, 0.0]
Enzyme.gradient(Enzyme.Reverse, Enzyme.Const(f), x)

Output

ERROR: LoadError: AssertionError: RT <: Array
Stacktrace:
  [1] abs_typeof(arg::LLVM.CallInst, partial::Bool)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/TiboG/src/absint.jl:266
  [2] abs_typeof
    @ ~/.julia/packages/Enzyme/TiboG/src/absint.jl:153 [inlined]
  [3] codegen(output::Symbol, job::GPUCompiler.CompilerJob{…}; libraries::Bool, deferred_codegen::Bool, optimize::Bool, toplevel::Bool, strip::Bool, validate::Bool, only_entry::Bool, parent_job::Nothing)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:6245
  [4] codegen
    @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:5614 [inlined]
  [5] _thunk(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams}, postopt::Bool)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:7241
  [6] _thunk
    @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:7241 [inlined]
  [7] cached_compilation
    @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:7282 [inlined]
  [8] thunkbase(ctx::LLVM.Context, mi::Core.MethodInstance, ::Val{…}, ::Type{…}, ::Type{…}, tt::Type{…}, ::Val{…}, ::Val{…}, ::Val{…}, ::Val{…}, ::Val{…}, ::Type{…}, ::Val{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:7355
  [9] #s2080#19052
    @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:7407 [inlined]
 [10]
    @ Enzyme.Compiler ./none:0
 [11] (::Core.GeneratedFunctionStub)(::UInt64, ::LineNumberNode, ::Any, ::Vararg{Any})
    @ Core ./boot.jl:602
 [12] autodiff
    @ ~/.julia/packages/Enzyme/TiboG/src/Enzyme.jl:315 [inlined]
 [13] gradient(rm::EnzymeCore.ReverseMode{…}, f::EnzymeCore.Const{…}, x::Vector{…})
    @ Enzyme ~/.julia/packages/Enzyme/TiboG/src/Enzyme.jl:1049
 [14] top-level scope
    @ ~/projects/Enzyme-mwes/RT_not_an_Array/mwe.jl:20
 [15] include(fname::String)
    @ Base.MainInclude ./client.jl:489
 [16] top-level scope
    @ REPL[11]:1
in expression starting at /Users/mhauru/projects/Enzyme-mwes/RT_not_an_Array/mwe.jl:1
Some type information was truncated. Use `show(err)` to see complete types.

[wsmoses]

@mhauru having a MWE for this as well would be quite helpful

[mhauru]

using Enzyme: Enzyme
using LinearAlgebra: Diagonal, transpose

function invwsumsq(w::AbstractVector, a::AbstractVector)
    s = zero(zero(eltype(a)) / zero(eltype(w)))
    for i in eachindex(w)
        s += abs2(a[i]) / w[i]
    end
    return s
end

_logpdf(d, x) = invwsumsq(d.Σ.diag, x .- d.μ)

function demo_func(x::Any=transpose([1.5 2.0;]);)
    m = [-0.30725218207431315, 0.5492115788562757]
    d = (; Σ = Diagonal([1.0, 1.0]), μ = m)
    logp = _logpdf(d, reshape(x, (2,)))
    return logp
end

f(x) = demo_func()

x = [0.0, 0.0]
Enzyme.autodiff(
    Enzyme.Reverse,
    Enzyme.Const(f),
    Enzyme.Active,
    Enzyme.Duplicated(x, zero(x)),
)

There you go.

[wsmoses]

Fixed by https://github.com/EnzymeAD/Enzyme.jl/pull/1884