Forward mode gradient AssertionError

[mhauru]

MWE:

using Turing: Turing
using Enzyme: Enzyme

Enzyme.API.runtimeActivity!(true)

Turing.@model function demo()
    x ~ Turing.Normal()
    y ~ Turing.Normal()
end

model = demo()
ctx = Turing.DefaultContext()
vi = Turing.VarInfo(model)
ldp = Turing.LogDensityFunction(vi, model, ctx)
x = randn(2)
f = Base.Fix1(Turing.LogDensityProblems.logdensity, ldp)
Enzyme.gradient(Enzyme.Forward, Enzyme.Const(f), x)

Output:

ERROR: LoadError: AssertionError: width == 1
Stacktrace:
  [1]
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:4277
  [2] enzyme!(job::GPUCompiler.CompilerJob{…}, mod::LLVM.Module, primalf::LLVM.Function, TT::Type, mode::Enzyme.API.CDerivativeMode, width::Int64, parallel::Bool, actualRetType::Type, wrap::Bool, modifiedBetween::Tuple{…}, returnPrimal::Bool, expectedTapeType::Type, loweredArgs::Set{…}, boxedArgs::Set{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:4188
  [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:6438
  [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] runtime_generic_fwd(activity::Type{…}, width::Val{…}, RT::Val{…}, f::Base.Fix1{…}, df::Base.Fix1{…}, df_2::Base.Fix1{…}, primal_1::AbstractPPL.VarName{…}, shadow_1_1::Nothing, shadow_1_2::Nothing)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/TiboG/src/rules/jitrules.jl:221
 [13] _all
    @ ./reduce.jl:1288
 [14] all
    @ ./reducedim.jl:1023 [inlined]
 [15] istrans
    @ ~/.julia/packages/DynamicPPL/DvdZw/src/abstract_varinfo.jl:531 [inlined]
 [16] istrans
    @ ~/.julia/packages/DynamicPPL/DvdZw/src/abstract_varinfo.jl:529
 [17] maybe_invlink_before_eval!!
    @ ~/.julia/packages/DynamicPPL/DvdZw/src/varinfo.jl:1071 [inlined]
 [18] macro expansion
    @ ~/.julia/packages/DynamicPPL/DvdZw/src/model.jl:1003 [inlined]
 [19] make_evaluate_args_and_kwargs
    @ ~/.julia/packages/DynamicPPL/DvdZw/src/model.jl:981 [inlined]
 [20] _evaluate!!
    @ ~/.julia/packages/DynamicPPL/DvdZw/src/model.jl:972 [inlined]
 [21] evaluate_threadunsafe!!
    @ ~/.julia/packages/DynamicPPL/DvdZw/src/model.jl:946 [inlined]
 [22] evaluate!!
    @ ~/.julia/packages/DynamicPPL/DvdZw/src/model.jl:894 [inlined]
 [23] logdensity
    @ ~/.julia/packages/DynamicPPL/DvdZw/src/logdensityfunction.jl:140
 [24] Fix1
    @ ./operators.jl:1118 [inlined]
 [25] Fix1
    @ ./operators.jl:0 [inlined]
 [26] fwddiffe2julia_Fix1_16584_inner_1wrap
    @ ./operators.jl:0
 [27] macro expansion
    @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:7187 [inlined]
 [28] enzyme_call
    @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:6794 [inlined]
 [29] ForwardModeThunk
    @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:6674 [inlined]
 [30] autodiff
    @ ~/.julia/packages/Enzyme/TiboG/src/Enzyme.jl:437 [inlined]
 [31] gradient(::EnzymeCore.ForwardMode{…}, f::EnzymeCore.Const{…}, x::Vector{…}; shadow::Tuple{…})
    @ Enzyme ~/.julia/packages/Enzyme/TiboG/src/Enzyme.jl:1136
 [32] gradient(::EnzymeCore.ForwardMode{EnzymeCore.FFIABI, false}, f::EnzymeCore.Const{Base.Fix1{…}}, x::Vector{Float64})
    @ Enzyme ~/.julia/packages/Enzyme/TiboG/src/Enzyme.jl:1132
 [33] top-level scope
    @ ~/projects/Enzyme-mwes/width-assertion-error/mwe.jl:19
 [34] include(fname::String)
    @ Base.MainInclude ./client.jl:489
 [35] top-level scope
    @ REPL[3]:1
in expression starting at /Users/mhauru/projects/Enzyme-mwes/width-assertion-error/mwe.jl:1
Some type information was truncated. Use `show(err)` to see complete types.

[wsmoses]

This merits investigation, but also TuringLang/DynamicPPL.jl#658 may fix it -- give it a shot?

[mhauru]

That does indeed make the issue go away.

[wsmoses]

Should be resolved separately by #1835