Hessian of a vector valued function

[vchuravy]

I was playing around with JuliaDiff/ForwardDiff.jl#61 (comment) to calculate the hessian of a vector-valued function.

function hessian(f::F, x) where F
    n = length(f(x))
    out = Enzyme.jacobian(Forward, x -> Enzyme.jacobian(Forward, f, x), x)
    return reshape(out, (n, eachindex(x), eachindex(x)))
end

function F!(res, x)
    res[1] = x[1]^2 + x[2]^2 - 2
    return res[2] = exp(x[1] - 1) + x[2]^2 - 2
end

function F(x)
    res = similar(x)
    F!(res, x)
    return res
end

julia> hessian(F, [3.0, 5.0])
ERROR: AssertionError: ; Function Attrs: mustprogress nofree willreturn
define internal fastcc nonnull "enzyme_type"="{[-1]:Pointer}" [2 x {} addrspace(10)*] @fwddiffe2julia_gradient_6822({} addrspace(10)* nofree noundef nonnull align 16 dereferenceable(40) "enzyme_type"="{[-1]:Pointer, [-1,0]:Pointer, [-1,0,-1]:Float@double, [-1,8]:Integer, [-1,9]:Integer, [-1,10]:Integer, [-1,11]:Integer, [-1,12]:Integer, [-1,13]:Integer, [-1,14]:Integer, [-1,15]:Integer, [-1,16]:Integer, [-1,17]:Integer, [-1,18]:Integer, [-1,19]:Integer, [-1,20]:Integer, [-1,21]:Integer, [-1,22]:Integer, [-1,23]:Integer, [-1,24]:Integer, [-1,25]:Integer, [-1,26]:Integer, [-1,27]:Integer, [-1,28]:Integer, [-1,29]:Integer, [-1,30]:Integer, [-1,31]:Integer, [-1,32]:Integer, [-1,33]:Integer, [-1,34]:Integer, [-1,35]:Integer, [-1,36]:Integer, [-1,37]:Integer, [-1,38]:Integer, [-1,39]:Integer}" "enzymejl_parmtype"="127053186628544" "enzymejl_parmtype_ref"="2" %0, [2 x {} addrspace(10)*] %"'") unnamed_addr #14 !dbg !157 {
top:
  %1 = call {}*** @julia.get_pgcstack() #15
  %current_task16 = getelementptr inbounds {}**, {}*** %1, i64 -14
  %current_task1 = bitcast {}*** %current_task16 to {}**
  %ptls_field7 = getelementptr inbounds {}**, {}*** %1, i64 2
  %2 = bitcast {}*** %ptls_field7 to i64***
  %ptls_load89 = load i64**, i64*** %2, align 8, !tbaa !10, !alias.scope !158, !noalias !161
  %3 = getelementptr inbounds i64*, i64** %ptls_load89, i64 2
  %safepoint = load i64*, i64** %3, align 8, !tbaa !14, !alias.scope !164, !noalias !167
  fence syncscope("singlethread") seq_cst
  call void @julia.safepoint(i64* %safepoint) #15, !dbg !170
  fence syncscope("singlethread") seq_cst
  %4 = addrspacecast {} addrspace(10)* %0 to { i8 addrspace(13)*, i64, i16, i16, i32 } addrspace(11)*, !dbg !171
  %arraylen_ptr = getelementptr inbounds { i8 addrspace(13)*, i64, i16, i16, i32 }, { i8 addrspace(13)*, i64, i16, i16, i32 } addrspace(11)* %4, i64 0, i32 1, !dbg !171
  %arraylen = load i64, i64 addrspace(11)* %arraylen_ptr, align 8, !dbg !171, !tbaa !26, !range !29, !alias.scope !175, !noalias !178, !enzyme_type !38, !enzymejl_source_type_UInt64 !9, !enzymejl_byref_BITS_VALUE !9, !enzyme_inactive !9
  %5 = call noalias "enzyme_inactive" "enzyme_type"="{[-1]:Pointer, [-1,-1]:Integer}" {} addrspace(10)* @ijl_box_int64(i64 %arraylen) #16, !dbg !181
  %6 = call "enzyme_inactive" "enzyme_type"="{[-1]:Pointer}" {} addrspace(10)* ({} addrspace(10)* ({} addrspace(10)*, {} addrspace(10)**, i32)*, {} addrspace(10)*, ...) @julia.call({} addrspace(10)* ({} addrspace(10)*, {} addrspace(10)**, i32)* noundef nonnull @jl_f_apply_type, {} addrspace(10)* noundef null, {} addrspace(10)* addrspacecast ({}* inttoptr (i64 127053190320704 to {}*) to {} addrspace(10)*), {} addrspace(10)* nonnull %5, {} addrspace(10)* addrspacecast ({}* inttoptr (i64 127053186628544 to {}*) to {} addrspace(10)*)) #17, !dbg !181
  %7 = shl nuw i64 %arraylen, 3, !dbg !181
  %8 = call noalias "enzyme_type"="{[-1]:Pointer}" {} addrspace(10)* @julia.gc_alloc_obj({}** nonnull %current_task1, i64 %7, {} addrspace(10)* nonnull %6) #18, !dbg !181
  %9 = call noalias "enzyme_type"="{[-1]:Pointer}" {} addrspace(10)* @julia.gc_alloc_obj({}** nonnull %current_task1, i64 %7, {} addrspace(10)* nonnull %6) #18, !dbg !181
  %"'dual_phi1" = phi [2 x {} addrspace(10)*] , !dbg !181
  %10 = bitcast {} addrspace(10)* %8 to i8 addrspace(10)*, !dbg !181
  %"'dual_phi2" = phi [2 x i8 addrspace(10)*] , !dbg !181
  call void @llvm.memset.p10i8.i64(i8 addrspace(10)* nonnull align 1 %10, i8 noundef 0, i64 %7, i1 noundef false) #15, !dbg !181
  %11 = bitcast {} addrspace(10)* %8 to {} addrspace(10)* addrspace(10)*, !dbg !181
  %"'dual_phi3" = phi [2 x {} addrspace(10)* addrspace(10)*] , !dbg !181
  %12 = addrspacecast {} addrspace(10)* addrspace(10)* %11 to {} addrspace(10)* addrspace(11)*, !dbg !181
  %"'dual_phi4" = phi [2 x {} addrspace(10)* addrspace(11)*] , !dbg !181
  %13 = icmp eq i64 %arraylen, 0, !dbg !181
  br i1 %13, label %julia_gradient_6822u6824.exit, label %loop.i.preheader, !dbg !181

loop.i.preheader:                                 ; preds = %top
  %14 = add nsw i64 %arraylen, -1, !dbg !183
  br label %loop.i, !dbg !183

loop.i:                                           ; preds = %julia_zerosetfn_6431.exit.i, %loop.i.preheader
  %iv = phi i64 [ 0, %loop.i.preheader ], [ %iv.next, %julia_zerosetfn_6431.exit.i ], !dbg !181
  %iv.next = add nuw nsw i64 %iv, 1, !dbg !181
  %arraylen.i.i = load i64, i64 addrspace(11)* %arraylen_ptr, align 8, !dbg !187, !tbaa !26, !range !29, !alias.scope !30, !noalias !33, !enzyme_type !38, !enzymejl_source_type_UInt64 !9, !enzymejl_byref_BITS_VALUE !9, !enzyme_inactive !9
  %15 = call noalias nonnull "enzyme_type"="{[-1]:Pointer, [-1,0]:Pointer, [-1,0,-1]:Float@double, [-1,8]:Integer, [-1,9]:Integer, [-1,10]:Integer, [-1,11]:Integer, [-1,12]:Integer, [-1,13]:Integer, [-1,14]:Integer, [-1,15]:Integer, [-1,16]:Integer, [-1,17]:Integer, [-1,18]:Integer, [-1,19]:Integer, [-1,20]:Integer, [-1,21]:Integer, [-1,22]:Integer, [-1,23]:Integer, [-1,24]:Integer, [-1,25]:Integer, [-1,26]:Integer, [-1,27]:Integer, [-1,28]:Integer, [-1,29]:Integer, [-1,30]:Integer, [-1,31]:Integer, [-1,32]:Integer, [-1,33]:Integer, [-1,34]:Integer, [-1,35]:Integer, [-1,36]:Integer, [-1,37]:Integer, [-1,38]:Integer, [-1,39]:Integer}" {} addrspace(10)* @ijl_alloc_array_1d({} addrspace(10)* noundef addrspacecast ({}* inttoptr (i64 127053186628544 to {}*) to {} addrspace(10)*), i64 %arraylen.i.i) #16, !dbg !189
  %"'dual_phi5" = phi [2 x {} addrspace(10)*] , !dbg !190
  %16 = addrspacecast {} addrspace(10)* %15 to { i8 addrspace(13)*, i64, i16, i16, i32 } addrspace(11)*, !dbg !190
  %"'dual_phi6" = phi [2 x { i8 addrspace(13)*, i64, i16, i16, i32 } addrspace(11)*] , !dbg !190
  %arraylen_ptr2.i.i = getelementptr inbounds { i8 addrspace(13)*, i64, i16, i16, i32 }, { i8 addrspace(13)*, i64, i16, i16, i32 } addrspace(11)* %16, i64 0, i32 1, !dbg !190
  %"arraylen_ptr2.i.i'dual_phi" = phi [2 x i64 addrspace(11)*] , !dbg !190
  %arraylen3.i.i = load i64, i64 addrspace(11)* %arraylen_ptr2.i.i, align 8, !dbg !190, !tbaa !26, !range !29, !alias.scope !30, !noalias !33, !enzyme_type !38, !enzymejl_source_type_UInt64 !9, !enzymejl_byref_BITS_VALUE !9, !enzyme_inactive !9
  %.not = icmp eq i64 %arraylen3.i.i, 0, !dbg !195
  br i1 %.not, label %loop.i.julia_zerosetfn_6431.exit.i_crit_edge, label %loop.i.L15.i.i_crit_edge, !dbg !183

loop.i.julia_zerosetfn_6431.exit.i_crit_edge:     ; preds = %loop.i
  %17 = addrspacecast {} addrspace(10)* %15 to double addrspace(13)* addrspace(11)*, !dbg !199
  %"'dual_phi7" = phi [2 x double addrspace(13)* addrspace(11)*] , !dbg !199
  %arrayptr12.i.i.pre12 = load double addrspace(13)*, double addrspace(13)* addrspace(11)* %17, align 8, !dbg !199, !tbaa !81, !alias.scope !201, !noalias !33, !enzyme_type !86, !enzymejl_byref_BITS_VALUE !9, !enzymejl_source_type_Ptr\7BFloat64\7D !9
  %18 = bitcast {} addrspace(10)* %15 to double addrspace(13)* addrspace(10)*, !dbg !183
  %19 = bitcast double addrspace(13)* addrspace(10)* %18 to {} addrspace(10)*, !dbg !183
  br label %julia_zerosetfn_6431.exit.i, !dbg !183

loop.i.L15.i.i_crit_edge:                         ; preds = %loop.i
  %20 = addrspacecast {} addrspace(10)* %15 to double addrspace(13)* addrspace(11)*, !dbg !204
  %"'dual_phi8" = phi [2 x double addrspace(13)* addrspace(11)*] , !dbg !204
  %arrayptr.i.i.pre10 = load double addrspace(13)*, double addrspace(13)* addrspace(11)* %20, align 8, !dbg !204, !tbaa !81, !alias.scope !201, !noalias !33, !enzyme_type !86, !enzymejl_byref_BITS_VALUE !9, !enzymejl_source_type_Ptr\7BFloat64\7D !9
  %"arrayptr.i.i.pre10'dual_phi" = phi [2 x double addrspace(13)*] , !dbg !183
  %arrayptr.i.i.pre1013 = bitcast double addrspace(13)* %arrayptr.i.i.pre10 to i8 addrspace(13)*, !dbg !183
  %"arrayptr.i.i.pre1013'dual_phi" = phi [2 x i8 addrspace(13)*] , !dbg !183
  %21 = shl nuw i64 %arraylen3.i.i, 3, !dbg !183
  call void @llvm.memset.p13i8.i64(i8 addrspace(13)* align 8 %arrayptr.i.i.pre1013, i8 noundef 0, i64 %21, i1 noundef false) #15, !dbg !204, !tbaa !91, !alias.scope !94, !noalias !206
  %22 = bitcast {} addrspace(10)* %15 to double addrspace(13)* addrspace(10)*, !dbg !199
  %23 = bitcast double addrspace(13)* addrspace(10)* %22 to {} addrspace(10)*, !dbg !199
  br label %julia_zerosetfn_6431.exit.i, !dbg !199

julia_zerosetfn_6431.exit.i:                      ; preds = %loop.i.L15.i.i_crit_edge, %loop.i.julia_zerosetfn_6431.exit.i_crit_edge
  %24 = bitcast {} addrspace(10)* %15 to double addrspace(13)* addrspace(10)*, !dbg !199
  %"'dual_phi9" = phi [2 x double addrspace(13)* addrspace(10)*] , !dbg !199
  %25 = addrspacecast double addrspace(13)* addrspace(10)* %24 to double addrspace(13)* addrspace(11)*, !dbg !199
  %"'dual_phi10" = phi [2 x double addrspace(13)* addrspace(11)*] , !dbg !199
  %26 = load double addrspace(13)*, double addrspace(13)* addrspace(11)* %25, align 8, !dbg !199, !enzyme_type !86, !enzymejl_byref_BITS_VALUE !9, !enzymejl_source_type_Ptr\7BFloat64\7D !9
  %"'dual_phi11" = phi [2 x double addrspace(13)*] , !dbg !199
  %27 = getelementptr inbounds double, double addrspace(13)* %26, i64 %iv, !dbg !199
  %"'dual_phi12" = phi [2 x double addrspace(13)*] , !dbg !199
  store double 1.000000e+00, double addrspace(13)* %27, align 8, !dbg !199, !tbaa !91, !alias.scope !94, !noalias !206
  %28 = getelementptr {} addrspace(10)*, {} addrspace(10)* addrspace(11)* %12, i64 %iv, !dbg !181
  %"'dual_phi13" = phi [2 x {} addrspace(10)* addrspace(11)*] , !dbg !181
  store {} addrspace(10)* %15, {} addrspace(10)* addrspace(11)* %28, align 8, !dbg !181, !noalias !207
  call void ({} addrspace(10)*, ...) @julia.write_barrier({} addrspace(10)* nonnull %8, {} addrspace(10)* nonnull %15) #15, !dbg !181
  %29 = icmp eq i64 %iv.next, %arraylen, !dbg !181
  br i1 %29, label %julia_gradient_6822u6824.exit.loopexit, label %loop.i, !dbg !181

julia_gradient_6822u6824.exit.loopexit:           ; preds = %julia_zerosetfn_6431.exit.i
  br label %julia_gradient_6822u6824.exit, !dbg !173

julia_gradient_6822u6824.exit:                    ; preds = %julia_gradient_6822u6824.exit.loopexit, %top
  %30 = call noalias nonnull "enzyme_type"="{[-1]:Pointer}" {} addrspace(10)* ({} addrspace(10)* ({} addrspace(10)*, {} addrspace(10)**, i32)*, {} addrspace(10)*, ...) @julia.call({} addrspace(10)* ({} addrspace(10)*, {} addrspace(10)**, i32)* noundef nonnull @jl_f_tuple, {} addrspace(10)* noundef null, {} addrspace(10)* nonnull %8) #18, !dbg !173
  %"'dual_phi14" = phi [2 x {} addrspace(10)*] , !dbg !170
  %31 = call nonnull "enzyme_type"="{[-1]:Pointer}" {} addrspace(10)* ({} addrspace(10)* ({} addrspace(10)*, {} addrspace(10)**, i32)*, {} addrspace(10)*, ...) @julia.call({} addrspace(10)* ({} addrspace(10)*, {} addrspace(10)**, i32)* noundef nonnull @ijl_apply_generic, {} addrspace(10)* noundef addrspacecast ({}* inttoptr (i64 127050644672944 to {}*) to {} addrspace(10)*), {} addrspace(10)* addrspacecast ({}* inttoptr (i64 127053414268936 to {}*) to {} addrspace(10)*), {} addrspace(10)* nonnull %30, {} addrspace(10)* addrspacecast ({}* inttoptr (i64 127050644677280 to {}*) to {} addrspace(10)*), {} addrspace(10)* addrspacecast ({}* inttoptr (i64 127052829645376 to {}*) to {} addrspace(10)*), {} addrspace(10)* addrspacecast ({}* inttoptr (i64 127050932732816 to {}*) to {} addrspace(10)*), {} addrspace(10)* nonnull %0) #19, !dbg !170
  %"'dual_phi15" = phi [2 x {} addrspace(10)*] , !dbg !170
  ret {} addrspace(10)* %31, !dbg !170

allocsForInversion:                               ; No predecessors!
  %"iv'ac" = alloca i64, align 8
}

 Allocation could not have its type statically determined   %9 = call noalias "enzyme_type"="{[-1]:Pointer}" {} addrspace(10)* @julia.gc_alloc_obj({}** nonnull %current_task1, i64 %7, {} addrspace(10)* nonnull %6) #18, !dbg !56
Stacktrace:
  [1] shadow_alloc_rewrite(V::Ptr{…}, gutils::Ptr{…}, Orig::Ptr{…}, idx::UInt64, prev::Ptr{…}, used::UInt8)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/QsaeA/src/compiler.jl:671
  [2] EnzymeCreateForwardDiff(logic::Enzyme.Logic, todiff::LLVM.Function, retType::Enzyme.API.CDIFFE_TYPE, constant_args::Vector{…}, TA::Enzyme.TypeAnalysis, returnValue::Bool, mode::Enzyme.API.CDerivativeMode, runtimeActivity::Bool, width::Int64, additionalArg::Ptr{…}, typeInfo::Enzyme.FnTypeInfo, uncacheable_args::Vector{…})
    @ Enzyme.API ~/.julia/packages/Enzyme/QsaeA/src/api.jl:334
  [3] 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{…} where N, returnPrimal::Bool, expectedTapeType::Type, loweredArgs::Set{…}, boxedArgs::Set{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/QsaeA/src/compiler.jl:1745
  [4] 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/QsaeA/src/compiler.jl:4550
  [5] codegen
    @ ~/.julia/packages/Enzyme/QsaeA/src/compiler.jl:3353 [inlined]
  [6] _thunk(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams}, postopt::Bool)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/QsaeA/src/compiler.jl:5410
  [7] _thunk
    @ ~/.julia/packages/Enzyme/QsaeA/src/compiler.jl:5410 [inlined]
  [8] cached_compilation
    @ ~/.julia/packages/Enzyme/QsaeA/src/compiler.jl:5462 [inlined]
  [9] thunkbase(mi::Core.MethodInstance, World::UInt64, FA::Type{…}, A::Type{…}, TT::Type, Mode::Enzyme.API.CDerivativeMode, width::Int64, ModifiedBetween::Tuple{…} where N, ReturnPrimal::Bool, ShadowInit::Bool, ABI::Type, ErrIfFuncWritten::Bool, RuntimeActivity::Bool, edges::Vector{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/QsaeA/src/compiler.jl:5573
 [10] thunk_generator(world::UInt64, source::LineNumberNode, FA::Type, A::Type, TT::Type, Mode::Enzyme.API.CDerivativeMode, Width::Int64, ModifiedBetween::Tuple{…} where N, ReturnPrimal::Bool, ShadowInit::Bool, ABI::Type, ErrIfFuncWritten::Bool, RuntimeActivity::Bool, self::Any, fakeworld::Any, fa::Type, a::Type, tt::Type, mode::Type, width::Type, modifiedbetween::Type, returnprimal::Type, shadowinit::Type, abi::Type, erriffuncwritten::Type, runtimeactivity::Type)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/QsaeA/src/compiler.jl:5758
 [11] autodiff
    @ ~/.julia/packages/Enzyme/QsaeA/src/Enzyme.jl:640 [inlined]
 [12] autodiff
    @ ~/.julia/packages/Enzyme/QsaeA/src/Enzyme.jl:524 [inlined]
 [13] macro expansion
    @ ~/.julia/packages/Enzyme/QsaeA/src/sugar.jl:680 [inlined]
 [14] gradient(::ForwardMode{…}, ::var"#65#66"{…}, ::Vector{…}; chunk::Nothing, shadows::Tuple{…})
    @ Enzyme ~/.julia/packages/Enzyme/QsaeA/src/sugar.jl:582
 [15] gradient(::ForwardMode{false, FFIABI, false, false}, ::var"#65#66"{typeof(F)}, ::Vector{Float64})
    @ Enzyme ~/.julia/packages/Enzyme/QsaeA/src/sugar.jl:582
 [16] #jacobian#128
    @ ~/.julia/packages/Enzyme/QsaeA/src/sugar.jl:789 [inlined]
 [17] jacobian
    @ ~/.julia/packages/Enzyme/QsaeA/src/sugar.jl:788 [inlined]
 [18] hessian(f::typeof(F), x::Vector{Float64})
    @ Main ~/src/NewtonKrylov/test/runtests.jl:46
 [19] top-level scope
    @ REPL[5]:1
Some type information was truncated. Use `show(err)` to see complete types.