EnzymeAD / Enzyme.jl

Julia bindings for the Enzyme automatic differentiator
https://enzyme.mit.edu
MIT License
439 stars 62 forks source link

Forward-over-reverse: hyperbolic functions trigger `MethodError: no method matching is_concrete_tuple(::LLVM.LLVMDouble)` #1772

Closed danielwe closed 2 weeks ago

danielwe commented 3 weeks ago

Forward-over-reverse may result in a MethodError that seems to indicate that typetree is receiving an instance where it expects a type. In the MWE below, tanh can be replaced by any hyperbolic function (sinh, cosh, coth). No other function I've tried has triggered the issue.

using Enzyme

f(x) = sum(tanh, x)  # or sinh/cosh/coth

function df!(dx, x)
    make_zero!(dx)
    autodiff_deferred(Reverse, f, Active, Duplicated(x, dx))
    return nothing
end

function hvp!(hv, v, x)
    make_zero!(hv)
    autodiff(Forward, df!, Const, Duplicated(make_zero(x), hv), Duplicated(x, v))
    return nothing
end

x = [0.5]

# primal: sanity check
@show f(x)

# gradient: works
dx = make_zero(x)
df!(dx, x)
@show dx

# hvp: error
v = first(onehot(x))
hv = make_zero(v)
hvp!(hv, v, x)
@show hv

Output:

f(x) = 0.46211715726000974
dx = [0.7864477329659275]
ERROR: LoadError: MethodError: no method matching is_concrete_tuple(::LLVM.LLVMDouble)

Closest candidates are:
  is_concrete_tuple(::Type{T2}) where T2
   @ Enzyme ~/.julia/packages/Enzyme/Tb3Iu/src/utils.jl:11

Stacktrace:
  [1] typetree_inner(T::Any, ctx::LLVM.Context, dl::String, seen::IdDict{Any, Union{Nothing, Enzyme.TypeTree}})
    @ Enzyme ~/.julia/packages/Enzyme/Tb3Iu/src/typetree.jl:218
  [2] typetree(T::Any, ctx::LLVM.Context, dl::String, seen::IdDict{Any, Union{Nothing, Enzyme.TypeTree}})
    @ Enzyme ~/.julia/packages/Enzyme/Tb3Iu/src/typetree.jl:88
  [3] codegen(output::Symbol, job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.Enzyme
CompilerParams}; libraries::Bool, deferred_codegen::Bool, optimize::Bool, toplevel::Bool, strip::Bool, validate
::Bool, only_entry::Bool, parent_job::Nothing)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/Tb3Iu/src/compiler.jl:6244
  [4] codegen
    @ ~/.julia/packages/Enzyme/Tb3Iu/src/compiler.jl:5591 [inlined]
  [5] _thunk(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams}, 
postopt::Bool)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/Tb3Iu/src/compiler.jl:7205
  [6] _thunk
    @ ~/.julia/packages/Enzyme/Tb3Iu/src/compiler.jl:7205 [inlined]
  [7] cached_compilation
    @ ~/.julia/packages/Enzyme/Tb3Iu/src/compiler.jl:7246 [inlined]
  [8] thunkbase(ctx::LLVM.Context, mi::Core.MethodInstance, ::Val{0x0000000000007afd}, ::Type{Const{typeof(df!)
}}, ::Type{Const}, tt::Type{Tuple{Duplicated{Vector{Float64}}, Duplicated{Vector{Float64}}}}, ::Val{Enzyme.API.
DEM_ForwardMode}, ::Val{1}, ::Val{(false, false, false)}, ::Val{false}, ::Val{false}, ::Type{FFIABI}, ::Val{tru
e})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/Tb3Iu/src/compiler.jl:7319
  [9] #s2055#18999
    @ ~/.julia/packages/Enzyme/Tb3Iu/src/compiler.jl:7371 [inlined]
 [10] var"#s2055#18999"(FA::Any, A::Any, TT::Any, Mode::Any, ModifiedBetween::Any, width::Any, ReturnPrimal::An
y, ShadowInit::Any, World::Any, ABI::Any, ErrIfFuncWritten::Any, ::Any, ::Type, ::Type, ::Type, tt::Any, ::Type
, ::Type, ::Type, ::Type, ::Type, ::Type, ::Any)
    @ Enzyme.Compiler ./none:0
 [11] (::Core.GeneratedFunctionStub)(::UInt64, ::LineNumberNode, ::Any, ::Vararg{Any})
    @ Core ./boot.jl:602
 [12] autodiff
    @ ~/.julia/packages/Enzyme/Tb3Iu/src/Enzyme.jl:435 [inlined]
 [13] autodiff
    @ ~/.julia/packages/Enzyme/Tb3Iu/src/Enzyme.jl:332 [inlined]
 [14] hvp!(hv::Vector{Float64}, v::Vector{Float64}, x::Vector{Float64})
    @ Main ~/bugs/enzymehyperbolic.jl:13
 [15] top-level scope
    @ ~/bugs/enzymehyperbolic.jl:30
in expression starting at /home/daniel/bugs/enzymehyperbolic.jl:30
wsmoses commented 2 weeks ago

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