Issue with nested call of `jacobian()`

[tduretz]

Hi,
I have prepared a MWE which causes an AssertionError (see below).
The aim of the code is to compute the gradient of a vector, for this I use jacobian() .
The evaluation of the vector involves the solution of a non-linear problem, for that I also use jacobian().
Any help or advice would be welcome.

Full assertion error message

  ERROR: AssertionError: ; Function Attrs: mustprogress willreturn
define internal nonnull "enzyme_type"="{}" "enzymejl_parmtype"="4790243968" "enzymejl_parmtype_ref"="1" { {} addrspace(10)*, [4 x {} addrspace(10)*] } @fwddiffe4julia_gradient_21766_inner.1({} addrspace(10)* noundef nonnull align 8 dereferenceable(24) "enzyme_type"="{[-1]:Pointer, [-1,0]:Pointer, [-1,0,-1]:Float@double, [-1,8]:Pointer, [-1,8,0]:Integer, [-1,8,1]:Integer, [-1,8,2]:Integer, [-1,8,3]:Integer, [-1,8,4]:Integer, [-1,8,5]:Integer, [-1,8,6]:Integer, [-1,8,7]:Integer, [-1,8,8]:Pointer, [-1,8,8,-1]:Float@double, [-1,16]:Integer, [-1,17]:Integer, [-1,18]:Integer, [-1,19]:Integer, [-1,20]:Integer, [-1,21]:Integer, [-1,22]:Integer, [-1,23]:Integer}" "enzymejl_parmtype"="4735203216" "enzymejl_parmtype_ref"="2" %0, [4 x {} addrspace(10)*] %"'", [1 x double] "enzyme_type"="{[-1]:Float@double}" "enzymejl_parmtype"="4533454096" "enzymejl_parmtype_ref"="0" %1, [1 x double] "enzyme_type"="{[-1]:Float@double}" "enzymejl_parmtype"="4533454096" "enzymejl_parmtype_ref"="0" %2, [4 x [1 x double]] %"'1", [1 x [3 x double]] "enzyme_type"="{[-1]:Float@double}" "enzymejl_parmtype"="13430207248" "enzymejl_parmtype_ref"="0" %3, [4 x [1 x [3 x double]]] %"'2") local_unnamed_addr #12 !dbg !152 {
entry:
  %pgcstack.i = call {}*** @julia.get_pgcstack() #13
  %4 = bitcast {} addrspace(10)* %0 to i8 addrspace(10)*, !dbg !153
  %5 = addrspacecast i8 addrspace(10)* %4 to i8 addrspace(11)*, !dbg !153
  %6 = getelementptr inbounds i8, i8 addrspace(11)* %5, i64 16, !dbg !153
  %7 = bitcast i8 addrspace(11)* %6 to i64 addrspace(11)*, !dbg !153
  %8 = load i64, i64 addrspace(11)* %7, align 8, !dbg !153, !tbaa !21, !alias.scope !159, !noalias !162, !enzyme_inactive !0, !enzyme_type !34, !enzymejl_source_type_Int64 !0, !enzymejl_byref_BITS_VALUE !0
  %9 = call noalias "enzyme_inactive" "enzyme_type"="{[-1]:Pointer, [-1,-1]:Integer}" {} addrspace(10)* @ijl_box_int64(i64 %8) #14, !dbg !167
  %10 = 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 4757086928 to {}*) to {} addrspace(10)*), {} addrspace(10)* nonnull %9, {} addrspace(10)* addrspacecast ({}* inttoptr (i64 4735203216 to {}*) to {} addrspace(10)*)) #15, !dbg !167
  %11 = shl nuw i64 %8, 3, !dbg !167
  %12 = getelementptr inbounds {}**, {}*** %pgcstack.i, i64 -14, !dbg !167
  %13 = bitcast {}*** %12 to {}*, !dbg !167
  %14 = call noalias "enzyme_type"="{[-1]:Pointer}" {} addrspace(10)* @julia.gc_alloc_obj({}* nonnull %13, i64 %11, {} addrspace(10)* nonnull %10) #16, !dbg !167
  %15 = call noalias "enzyme_type"="{[-1]:Pointer}" {} addrspace(10)* @julia.gc_alloc_obj({}* nonnull %13, i64 %11, {} addrspace(10)* nonnull %10) #16, !dbg !167
  %"'dual_phi6" = phi [4 x {} addrspace(10)*] , !dbg !167
  %16 = bitcast {} addrspace(10)* %14 to i8 addrspace(10)*, !dbg !167
  %"'dual_phi7" = phi [4 x i8 addrspace(10)*] , !dbg !167
  call void @llvm.memset.p10i8.i64(i8 addrspace(10)* nonnull align 1 %16, i8 noundef 0, i64 %11, i1 noundef false) #13, !dbg !167
  %17 = bitcast {} addrspace(10)* %14 to {} addrspace(10)* addrspace(10)*, !dbg !167
  %"'dual_phi8" = phi [4 x {} addrspace(10)* addrspace(10)*] , !dbg !167
  %18 = addrspacecast {} addrspace(10)* addrspace(10)* %17 to {} addrspace(10)* addrspace(11)*, !dbg !167
  %"'dual_phi9" = phi [4 x {} addrspace(10)* addrspace(11)*] , !dbg !167
  %19 = icmp eq i64 %8, 0, !dbg !167
  br i1 %19, label %julia_gradient_21766_inner.exit, label %loop.i.i.preheader, !dbg !167

loop.i.i.preheader:                               ; preds = %entry
  %20 = add i64 %8, -1, !dbg !169
  br label %loop.i.i, !dbg !169

loop.i.i:                                         ; preds = %julia_zerosetfn_14123.exit.i.i, %loop.i.i.preheader
  %iv = phi i64 [ 0, %loop.i.i.preheader ], [ %iv.next, %julia_zerosetfn_14123.exit.i.i ], !dbg !167
  %iv.next = add nuw nsw i64 %iv, 1, !dbg !167
  %21 = load i64, i64 addrspace(11)* %7, align 8, !dbg !175, !tbaa !21, !alias.scope !26, !noalias !29, !enzyme_inactive !0, !enzyme_type !34, !enzymejl_source_type_Int64 !0, !enzymejl_byref_BITS_VALUE !0
  %.not = icmp eq i64 %21, 0, !dbg !169
  br i1 %.not, label %L5.i.i.i, label %L7.i.i.i, !dbg !169

L5.i.i.i:                                         ; preds = %loop.i.i
  %22 = load atomic {} addrspace(10)*, {} addrspace(10)** inttoptr (i64 4735203440 to {} addrspace(10)**) unordered, align 16, !dbg !176, !tbaa !56, !alias.scope !58, !noalias !59, !nonnull !0, !enzyme_type !60, !enzymejl_source_type_Memory\7BFloat64\7D !0, !enzymejl_byref_BITS_REF !0
  %.not47 = icmp eq {} addrspace(10)* %22, null, !dbg !176
  br i1 %.not47, label %fail.i.i.i, label %L9.i.i.i, !dbg !176

L7.i.i.i:                                         ; preds = %loop.i.i
  %23 = call noalias "enzyme_type"="{[-1]:Pointer, [-1,0]:Integer, [-1,1]:Integer, [-1,2]:Integer, [-1,3]:Integer, [-1,4]:Integer, [-1,5]:Integer, [-1,6]:Integer, [-1,7]:Integer, [-1,8]:Pointer, [-1,8,-1]:Float@double}" {} addrspace(10)* @jl_alloc_genericmemory({} addrspace(10)* noundef addrspacecast ({}* inttoptr (i64 4735203408 to {}*) to {} addrspace(10)*), i64 %21) #14, !dbg !177
  %"'dual_phi10" = phi [4 x {} addrspace(10)*] , !dbg !177
  br label %L9.i.i.i, !dbg !177

L9.i.i.i:                                         ; preds = %L7.i.i.i, %L5.i.i.i
  %.pre41 = phi {} addrspace(10)* [ %23, %L7.i.i.i ], [ %22, %L5.i.i.i ], !dbg !167
  %".pre41'dual_phi" = phi [4 x {} addrspace(10)*] , !dbg !178
  %24 = bitcast {} addrspace(10)* %.pre41 to { i64, {} addrspace(10)** } addrspace(10)*, !dbg !178
  %"'dual_phi11" = phi [4 x { i64, {} addrspace(10)** } addrspace(10)*] , !dbg !178
  %25 = addrspacecast { i64, {} addrspace(10)** } addrspace(10)* %24 to { i64, {} addrspace(10)** } addrspace(11)*, !dbg !178
  %"'dual_phi12" = phi [4 x { i64, {} addrspace(10)** } addrspace(11)*] , !dbg !178
  %26 = getelementptr inbounds { i64, {} addrspace(10)** }, { i64, {} addrspace(10)** } addrspace(11)* %25, i64 0, i32 1, !dbg !178
  %"'dual_phi13" = phi [4 x {} addrspace(10)** addrspace(11)*] , !dbg !178
  %27 = bitcast {} addrspace(10)** addrspace(11)* %26 to i8* addrspace(11)*, !dbg !178
  %"'dual_phi14" = phi [4 x i8* addrspace(11)*] , !dbg !178
  %28 = load i8*, i8* addrspace(11)* %27, align 8, !dbg !178, !tbaa !56, !alias.scope !58, !noalias !59, !nonnull !0, !enzyme_type !68, !enzymejl_byref_BITS_VALUE !0, !enzymejl_source_type_Ptr\7BFloat64\7D !0, !enzyme_nocache !0
  %"'dual_phi15" = phi [4 x i8*] , !dbg !180
  %29 = call noalias nonnull align 8 dereferenceable(24) "enzyme_type"="{[-1]:Pointer, [-1,0]:Pointer, [-1,0,-1]:Float@double, [-1,8]:Pointer, [-1,8,0]:Integer, [-1,8,1]:Integer, [-1,8,2]:Integer, [-1,8,3]:Integer, [-1,8,4]:Integer, [-1,8,5]:Integer, [-1,8,6]:Integer, [-1,8,7]:Integer, [-1,8,8]:Pointer, [-1,8,8,-1]:Float@double, [-1,16]:Integer, [-1,17]:Integer, [-1,18]:Integer, [-1,19]:Integer, [-1,20]:Integer, [-1,21]:Integer, [-1,22]:Integer, [-1,23]:Integer}" {} addrspace(10)* @julia.gc_alloc_obj({}* nonnull %13, i64 noundef 24, {} addrspace(10)* noundef addrspacecast ({}* inttoptr (i64 4735203216 to {}*) to {} addrspace(10)*)) #17, !dbg !180
  %"'dual_phi16" = phi [4 x {} addrspace(10)*] , !dbg !180
  %30 = bitcast {} addrspace(10)* %29 to { i8*, {} addrspace(10)* } addrspace(10)*, !dbg !180
  %"'dual_phi17" = phi [4 x { i8*, {} addrspace(10)* } addrspace(10)*] , !dbg !180
  %31 = addrspacecast { i8*, {} addrspace(10)* } addrspace(10)* %30 to { i8*, {} addrspace(10)* } addrspace(11)*, !dbg !180
  %"'dual_phi18" = phi [4 x { i8*, {} addrspace(10)* } addrspace(11)*] , !dbg !180
  %.repack = getelementptr inbounds { i8*, {} addrspace(10)* }, { i8*, {} addrspace(10)* } addrspace(11)* %31, i64 0, i32 0, !dbg !180
  %".repack'dual_phi" = phi [4 x i8* addrspace(11)*] , !dbg !180
  store i8* %28, i8* addrspace(11)* %.repack, align 8, !dbg !180, !tbaa !70, !alias.scope !26, !noalias !181
  %.repack48 = getelementptr inbounds { i8*, {} addrspace(10)* }, { i8*, {} addrspace(10)* } addrspace(11)* %31, i64 0, i32 1, !dbg !180
  %".repack48'dual_phi" = phi [4 x {} addrspace(10)* addrspace(11)*] , !dbg !180
  store {} addrspace(10)* %.pre41, {} addrspace(10)* addrspace(11)* %.repack48, align 8, !dbg !180, !tbaa !70, !alias.scope !26, !noalias !181
  %32 = bitcast {} addrspace(10)* %29 to i8 addrspace(10)*, !dbg !180
  %"'dual_phi19" = phi [4 x i8 addrspace(10)*] , !dbg !180
  %33 = addrspacecast i8 addrspace(10)* %32 to i8 addrspace(11)*, !dbg !180
  %"'dual_phi20" = phi [4 x i8 addrspace(11)*] , !dbg !180
  %34 = getelementptr inbounds i8, i8 addrspace(11)* %33, i64 16, !dbg !180
  %"'dual_phi21" = phi [4 x i8 addrspace(11)*] , !dbg !180
  %35 = bitcast i8 addrspace(11)* %34 to i64 addrspace(11)*, !dbg !180
  %"'dual_phi22" = phi [4 x i64 addrspace(11)*] , !dbg !180
  store i64 %21, i64 addrspace(11)* %35, align 8, !dbg !180, !tbaa !21, !alias.scope !26, !noalias !181
  %36 = icmp slt i64 %21, 1, !dbg !184
  %37 = bitcast i8* %28 to {} addrspace(10)**, !dbg !189
  %"'dual_phi23" = phi [4 x {} addrspace(10)**] , !dbg !189
  br i1 %36, label %L9.i.i.i.julia_zerosetfn_14123.exit.i.i_crit_edge, label %julia_zerosetfn_14123.exit.i.i.loopexit, !dbg !189

L9.i.i.i.julia_zerosetfn_14123.exit.i.i_crit_edge: ; preds = %L9.i.i.i
  %.pre43 = call "enzyme_type"="{[-1]:Pointer, [-1,-1]:Float@double}" {} addrspace(10)* addrspace(13)* @julia.gc_loaded({} addrspace(10)* noundef %.pre41, {} addrspace(10)** noundef %37) #13, !dbg !190
  %38 = bitcast { i64, {} addrspace(10)** } addrspace(10)* %24 to {} addrspace(10)** addrspace(10)*, !dbg !189
  %39 = bitcast {} addrspace(10)** addrspace(10)* %38 to i8* addrspace(10)*, !dbg !189
  %40 = bitcast i8* addrspace(10)* %39 to {} addrspace(10)* addrspace(10)*, !dbg !189
  br label %julia_zerosetfn_14123.exit.i.i, !dbg !189

fail.i.i.i:                                       ; preds = %L5.i.i.i
  %41 = load {}*, {}** @jl_undefref_exception, align 8, !dbg !176, !tbaa !56, !alias.scope !58, !noalias !59, !nonnull !0
  %42 = addrspacecast {}* %41 to {} addrspace(12)*, !dbg !176
  call void @ijl_throw({} addrspace(12)* %42) #18, !dbg !176
  unreachable, !dbg !176

julia_zerosetfn_14123.exit.i.i.loopexit:          ; preds = %L9.i.i.i
  %43 = call "enzyme_type"="{[-1]:Pointer, [-1,-1]:Float@double}" {} addrspace(10)* addrspace(13)* @julia.gc_loaded({} addrspace(10)* noundef %.pre41, {} addrspace(10)** noundef %37) #13, !dbg !192
  %"'dual_phi24" = phi [4 x {} addrspace(10)* addrspace(13)*] , !dbg !189
  %44 = bitcast {} addrspace(10)* addrspace(13)* %43 to i8 addrspace(13)*, !dbg !189
  %"'dual_phi25" = phi [4 x i8 addrspace(13)*] , !dbg !189
  %45 = shl nuw i64 %21, 3, !dbg !189
  call void @llvm.memset.p13i8.i64(i8 addrspace(13)* noundef nonnull align 8 %44, i8 noundef 0, i64 %45, i1 noundef false) #13, !dbg !193, !tbaa !95, !alias.scope !98, !noalias !195
  %46 = bitcast { i64, {} addrspace(10)** } addrspace(10)* %24 to {} addrspace(10)** addrspace(10)*, !dbg !190
  %47 = bitcast {} addrspace(10)** addrspace(10)* %46 to i8* addrspace(10)*, !dbg !190
  %48 = bitcast i8* addrspace(10)* %47 to {} addrspace(10)* addrspace(10)*, !dbg !190
  br label %julia_zerosetfn_14123.exit.i.i, !dbg !190

julia_zerosetfn_14123.exit.i.i:                   ; preds = %julia_zerosetfn_14123.exit.i.i.loopexit, %L9.i.i.i.julia_zerosetfn_14123.exit.i.i_crit_edge
  %49 = bitcast {} addrspace(10)* %.pre41 to <{ i64, i8* }> addrspace(10)*, !dbg !190
  %"'dual_phi26" = phi [4 x <{ i64, i8* }> addrspace(10)*] , !dbg !190
  %50 = getelementptr inbounds <{ i64, i8* }>, <{ i64, i8* }> addrspace(10)* %49, i32 0, i32 1, !dbg !190
  %"'dual_phi27" = phi [4 x i8* addrspace(10)*] , !dbg !190
  %51 = load i8*, i8* addrspace(10)* %50, align 8, !dbg !190, !enzyme_type !68, !enzymejl_byref_BITS_VALUE !0, !enzymejl_source_type_Ptr\7BFloat64\7D !0, !enzyme_nocache !0
  %"'dual_phi28" = phi [4 x i8*] , !dbg !190
  %52 = bitcast i8* %51 to {} addrspace(10)**, !dbg !190
  %"'dual_phi29" = phi [4 x {} addrspace(10)**] , !dbg !190
  %53 = call "enzyme_type"="{[-1]:Pointer, [-1,-1]:Float@double}" {} addrspace(10)* addrspace(13)* @julia.gc_loaded({} addrspace(10)* %.pre41, {} addrspace(10)** %52) #13, !dbg !190
  %"'dual_phi30" = phi [4 x {} addrspace(10)* addrspace(13)*] , !dbg !190
  %54 = getelementptr inbounds {} addrspace(10)*, {} addrspace(10)* addrspace(13)* %53, i64 %iv, !dbg !190
  %"'dual_phi31" = phi [4 x {} addrspace(10)* addrspace(13)*] , !dbg !190
  %55 = bitcast {} addrspace(10)* addrspace(13)* %54 to double addrspace(13)*, !dbg !190
  %"'dual_phi32" = phi [4 x double addrspace(13)*] , !dbg !190
  store double 1.000000e+00, double addrspace(13)* %55, align 8, !dbg !190, !tbaa !95, !alias.scope !98, !noalias !195
  %56 = getelementptr {} addrspace(10)*, {} addrspace(10)* addrspace(11)* %18, i64 %iv, !dbg !167
  %"'dual_phi33" = phi [4 x {} addrspace(10)* addrspace(11)*] , !dbg !167
  store {} addrspace(10)* %29, {} addrspace(10)* addrspace(11)* %56, align 8, !dbg !167, !noalias !196
  call void ({} addrspace(10)*, ...) @julia.write_barrier({} addrspace(10)* nonnull %14, {} addrspace(10)* nofree nonnull %29) #13, !dbg !167
  %57 = icmp eq i64 %iv.next, %8, !dbg !167
  br i1 %57, label %julia_gradient_21766_inner.exit.loopexit, label %loop.i.i, !dbg !167

julia_gradient_21766_inner.exit.loopexit:         ; preds = %julia_zerosetfn_14123.exit.i.i
  br label %julia_gradient_21766_inner.exit, !dbg !192

julia_gradient_21766_inner.exit:                  ; preds = %julia_gradient_21766_inner.exit.loopexit, %entry
  %.fca.0.2.extract = extractvalue [1 x [3 x double]] %3, 0, 2, !dbg !192, !enzyme_type !101, !enzymejl_byref_BITS_VALUE !0, !enzymejl_source_type_Float64 !0
  %".fca.0.2.extract'dual_phi" = phi fast [4 x double] , !dbg !192
  %.fca.0.1.extract = extractvalue [1 x [3 x double]] %3, 0, 1, !dbg !192, !enzyme_type !101, !enzymejl_byref_BITS_VALUE !0, !enzymejl_source_type_Float64 !0
  %".fca.0.1.extract'dual_phi" = phi fast [4 x double] , !dbg !192
  %.fca.0.0.extract = extractvalue [1 x [3 x double]] %3, 0, 0, !dbg !192, !enzyme_type !101, !enzymejl_byref_BITS_VALUE !0, !enzymejl_source_type_Float64 !0
  %".fca.0.0.extract'dual_phi" = phi fast [4 x double] , !dbg !192
  %.fca.0.extract = extractvalue [1 x double] %2, 0, !dbg !192, !enzyme_type !101, !enzymejl_byref_BITS_VALUE !0, !enzymejl_source_type_Float64 !0
  %".fca.0.extract'dual_phi" = phi fast [4 x double] , !dbg !192
  %.fca.0.extract6 = extractvalue [1 x double] %1, 0, !dbg !192, !enzyme_type !101, !enzymejl_byref_BITS_VALUE !0, !enzymejl_source_type_Float64 !0
  %58 = load {}*, {}** @jl_nothing, align 8, !dbg !197, !tbaa !56, !alias.scope !58, !noalias !59, !nonnull !0, !enzyme_inactive !0, !enzyme_type !103, !enzymejl_byref_BITS_REF !0, !enzymejl_source_type_Nothing !0
  %59 = addrspacecast {}* %58 to {} addrspace(10)*, !dbg !197, !enzyme_inactive !0
  %60 = 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 %14, {} addrspace(10)* %59, {} addrspace(10)* %59, {} addrspace(10)* %59) #16, !dbg !197
  %"'dual_phi34" = phi [4 x {} addrspace(10)*] , !dbg !198
  %61 = call noalias nonnull align 8 dereferenceable(8) "enzyme_type"="{[-1]:Pointer, [-1,-1]:Float@double}" {} addrspace(10)* @julia.gc_alloc_obj({}* nonnull %13, i64 noundef 8, {} addrspace(10)* noundef addrspacecast ({}* inttoptr (i64 4533454096 to {}*) to {} addrspace(10)*)) #17, !dbg !198
  %"'dual_phi35" = phi [4 x {} addrspace(10)*] , !dbg !198
  %62 = bitcast {} addrspace(10)* %61 to double addrspace(10)*, !dbg !198
  %"'dual_phi36" = phi [4 x double addrspace(10)*] , !dbg !198
  store double %.fca.0.extract6, double addrspace(10)* %62, align 8, !dbg !198, !tbaa !106, !alias.scope !98, !noalias !195
  %63 = call noalias nonnull align 8 dereferenceable(8) "enzyme_type"="{[-1]:Pointer, [-1,-1]:Float@double}" {} addrspace(10)* @julia.gc_alloc_obj({}* nonnull %13, i64 noundef 8, {} addrspace(10)* noundef addrspacecast ({}* inttoptr (i64 4533454096 to {}*) to {} addrspace(10)*)) #17, !dbg !198
  %"'dual_phi37" = phi [4 x {} addrspace(10)*] , !dbg !198
  %64 = bitcast {} addrspace(10)* %63 to double addrspace(10)*, !dbg !198
  %"'dual_phi38" = phi [4 x double addrspace(10)*] , !dbg !198
  store double %.fca.0.extract, double addrspace(10)* %64, align 8, !dbg !198, !tbaa !106, !alias.scope !98, !noalias !195
  %65 = call noalias nonnull align 8 dereferenceable(24) "enzyme_type"="{[-1]:Pointer, [-1,-1]:Float@double}" {} addrspace(10)* @julia.gc_alloc_obj({}* nonnull %13, i64 noundef 24, {} addrspace(10)* noundef addrspacecast ({}* inttoptr (i64 13430207248 to {}*) to {} addrspace(10)*)) #17, !dbg !198
  %"'dual_phi39" = phi [4 x {} addrspace(10)*] , !dbg !198
  %66 = bitcast {} addrspace(10)* %65 to i8 addrspace(10)*, !dbg !198
  %"'dual_phi40" = phi [4 x i8 addrspace(10)*] , !dbg !198
  %args.i.sroa.6.16..sroa_cast = bitcast {} addrspace(10)* %65 to double addrspace(10)*, !dbg !198
  %"args.i.sroa.6.16..sroa_cast'dual_phi" = phi [4 x double addrspace(10)*] , !dbg !198
  store double %.fca.0.0.extract, double addrspace(10)* %args.i.sroa.6.16..sroa_cast, align 8, !dbg !198, !tbaa !109, !alias.scope !110, !noalias !199
  %args.i.sroa.8.16..sroa_idx = getelementptr inbounds i8, i8 addrspace(10)* %66, i64 8, !dbg !198
  %"args.i.sroa.8.16..sroa_idx'dual_phi" = phi [4 x i8 addrspace(10)*] , !dbg !198
  %args.i.sroa.8.16..sroa_cast = bitcast i8 addrspace(10)* %args.i.sroa.8.16..sroa_idx to double addrspace(10)*, !dbg !198
  %"args.i.sroa.8.16..sroa_cast'dual_phi" = phi [4 x double addrspace(10)*] , !dbg !198
  store double %.fca.0.1.extract, double addrspace(10)* %args.i.sroa.8.16..sroa_cast, align 8, !dbg !198, !tbaa !109, !alias.scope !110, !noalias !199
  %args.i.sroa.9.16..sroa_idx = getelementptr inbounds i8, i8 addrspace(10)* %66, i64 16, !dbg !198
  %"args.i.sroa.9.16..sroa_idx'dual_phi" = phi [4 x i8 addrspace(10)*] , !dbg !198
  %args.i.sroa.9.16..sroa_cast = bitcast i8 addrspace(10)* %args.i.sroa.9.16..sroa_idx to double addrspace(10)*, !dbg !198
  %"args.i.sroa.9.16..sroa_cast'dual_phi" = phi [4 x double addrspace(10)*] , !dbg !198
  store double %.fca.0.2.extract, double addrspace(10)* %args.i.sroa.9.16..sroa_cast, align 8, !dbg !198, !tbaa !109, !alias.scope !110, !noalias !199
  %67 = 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 13811727648 to {}*) to {} addrspace(10)*), {} addrspace(10)* %59, {} addrspace(10)* nonnull %60, {} addrspace(10)* addrspacecast ({}* inttoptr (i64 13811732352 to {}*) to {} addrspace(10)*), {} addrspace(10)* addrspacecast ({}* inttoptr (i64 4962778720 to {}*) to {} addrspace(10)*), {} addrspace(10)* addrspacecast ({}* inttoptr (i64 5221939448 to {}*) to {} addrspace(10)*), {} addrspace(10)* nofree nonnull %0, {} addrspace(10)* nofree nonnull %61, {} addrspace(10)* nofree nonnull %63, {} addrspace(10)* nofree nonnull %65) #19, !dbg !198
  %"'dual_phi41" = phi [4 x {} addrspace(10)*] , !dbg !192
  ret {} addrspace(10)* %67, !dbg !192

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

 Allocation could not have its type statically determined   %15 = call noalias "enzyme_type"="{[-1]:Pointer}" {} addrspace(10)* @julia.gc_alloc_obj({}* nonnull %13, i64 %11, {} addrspace(10)* nonnull %10) #16, !dbg !42
Stacktrace:
  [1] shadow_alloc_rewrite(V::Ptr{…}, gutils::Ptr{…}, Orig::Ptr{…}, idx::UInt64, prev::Ptr{…}, used::UInt8)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/NB8vY/src/compiler.jl:681
  [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/NB8vY/src/api.jl:338
  [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::NTuple{…} where N, returnPrimal::Bool, expectedTapeType::Type, loweredArgs::Set{…}, boxedArgs::Set{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/NB8vY/src/compiler.jl:1787
  [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/NB8vY/src/compiler.jl:4658
  [5] codegen
    @ ~/.julia/packages/Enzyme/NB8vY/src/compiler.jl:3444 [inlined]
  [6] _thunk(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams}, postopt::Bool)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/NB8vY/src/compiler.jl:5518
  [7] _thunk
    @ ~/.julia/packages/Enzyme/NB8vY/src/compiler.jl:5518 [inlined]
  [8] cached_compilation
    @ ~/.julia/packages/Enzyme/NB8vY/src/compiler.jl:5570 [inlined]
  [9] thunkbase(mi::Core.MethodInstance, World::UInt64, FA::Type{…}, A::Type{…}, TT::Type, Mode::Enzyme.API.CDerivativeMode, width::Int64, ModifiedBetween::NTuple{…} where N, ReturnPrimal::Bool, ShadowInit::Bool, ABI::Type, ErrIfFuncWritten::Bool, RuntimeActivity::Bool, edges::Vector{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/NB8vY/src/compiler.jl:5681
 [10] thunk_generator(world::UInt64, source::LineNumberNode, FA::Type, A::Type, TT::Type, Mode::Enzyme.API.CDerivativeMode, Width::Int64, ModifiedBetween::NTuple{…} 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/NB8vY/src/compiler.jl:5866
 [11] runtime_generic_fwd(activity::Type{…}, runtimeActivity::Val{…}, width::Val{…}, RT::Val{…}, f::typeof(gradient), df::Nothing, df_2::Nothing, df_3::Nothing, df_4::Nothing, primal_1::ForwardMode{…}, shadow_1_1::Nothing, shadow_1_2::Nothing, shadow_1_3::Nothing, shadow_1_4::Nothing, primal_2::typeof(residual_single_phase), shadow_2_1::Nothing, shadow_2_2::Nothing, shadow_2_3::Nothing, shadow_2_4::Nothing, primal_3::Vector{…}, shadow_3_1::Vector{…}, shadow_3_2::Vector{…}, shadow_3_3::Vector{…}, shadow_3_4::Vector{…}, primal_4::Const{…}, shadow_4_1::Nothing, shadow_4_2::Nothing, shadow_4_3::Nothing, shadow_4_4::Nothing, primal_5::Const{…}, shadow_5_1::Const{…}, shadow_5_2::Const{…}, shadow_5_3::Const{…}, shadow_5_4::Const{…}, primal_6::Const{…}, shadow_6_1::Const{…}, shadow_6_2::Const{…}, shadow_6_3::Const{…}, shadow_6_4::Const{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/NB8vY/src/rules/jitrules.jl:288
 [12] #jacobian#128
    @ ~/.julia/packages/Enzyme/NB8vY/src/sugar.jl:789 [inlined]
 [13] jacobian
    @ ~/.julia/packages/Enzyme/NB8vY/src/sugar.jl:788 [inlined]
 [14] StressVector
    @ ~/REPO/StagFDTools/standalone/SinglePhaseReturnMapping_Enzyme_MWE.jl:27 [inlined]
 [15] StressVector
    @ ~/REPO/StagFDTools/standalone/SinglePhaseReturnMapping_Enzyme_MWE.jl:0 [inlined]
 [16] fwddiffe4julia_StressVector_14149_inner_1wrap
    @ ~/REPO/StagFDTools/standalone/SinglePhaseReturnMapping_Enzyme_MWE.jl:0
 [17] macro expansion
    @ ~/.julia/packages/Enzyme/NB8vY/src/compiler.jl:5448 [inlined]
 [18] enzyme_call
    @ ~/.julia/packages/Enzyme/NB8vY/src/compiler.jl:4986 [inlined]
 [19] ForwardModeThunk
    @ ~/.julia/packages/Enzyme/NB8vY/src/compiler.jl:4874 [inlined]
 [20] autodiff
    @ ~/.julia/packages/Enzyme/NB8vY/src/Enzyme.jl:654 [inlined]
 [21] autodiff
    @ ~/.julia/packages/Enzyme/NB8vY/src/Enzyme.jl:524 [inlined]
 [22] macro expansion
    @ ~/.julia/packages/Enzyme/NB8vY/src/sugar.jl:680 [inlined]
 [23] gradient(fm::ForwardMode{…}, f::typeof(StressVector), x::Vector{…}, args::Const{…}; chunk::Nothing, shadows::Tuple{…})
    @ Enzyme ~/.julia/packages/Enzyme/NB8vY/src/sugar.jl:582
 [24] gradient(fm::ForwardMode{…}, f::typeof(StressVector), x::Vector{…}, args::Const{…})
    @ Enzyme ~/.julia/packages/Enzyme/NB8vY/src/sugar.jl:582
 [25] #jacobian#128
    @ ~/.julia/packages/Enzyme/NB8vY/src/sugar.jl:789 [inlined]
 [26] jacobian
    @ ~/.julia/packages/Enzyme/NB8vY/src/sugar.jl:788 [inlined]
 [27] single_phase_return_mapping()
    @ Main ~/REPO/StagFDTools/standalone/SinglePhaseReturnMapping_Enzyme_MWE.jl:56
 [28] top-level scope
    @ ~/REPO/StagFDTools/standalone/SinglePhaseReturnMapping_Enzyme_MWE.jl:60
Some type information was truncated. Use `show(err)` to see complete types.

`versioninfo()`

``` julia> versioninfo() Julia Version 1.11.4 Commit 8561cc3d68d (2025-03-10 11:36 UTC) Build Info: Official https://julialang.org/ release Platform Info: OS: macOS (arm64-apple-darwin24.0.0) CPU: 12 × Apple M2 Max WORD_SIZE: 64 LLVM: libLLVM-16.0.6 (ORCJIT, apple-m2) Threads: 8 default, 0 interactive, 4 GC (on 8 virtual cores) Environment: JULIA_EDITOR = code JULIA_VSCODE_REPL = 1 ```

`status Enzyme`

Project StagFDTools v0.1.0 Status `~/REPO/StagFDTools/Project.toml` ⌃ [7da242da] Enzyme v0.13.37

using Enzyme 
import LinearAlgebra: norm

# Function to be minimized via inner Newton iteration 
function residual_single_phase(x, ε̇II_eff, divV, p)
    G, K, C = p.G, p.K, p.C
    eps    = -1e-13
    f      = x[1]  - C
    return [ 
        ε̇II_eff    -  x[1]/G - x[3]*(f>=eps),
        divV       + (x[2])/K,
        f*(f>=eps) +  x[3]*(f<eps)
    ]
end

# Evaluate the stress vector involving the solution of a non-linear function
function StressVector(ϵ̇, params)

    ε̇_eff, divV = ϵ̇[1:3], ϵ̇[4]
    ε̇II_eff = 1.0

    # Initialise inner solution array
    x = [1.0, 1.0, 0.0]

    # Newton iteration
    for iter=1:10
        J = Enzyme.jacobian(Enzyme.ForwardWithPrimal, residual_single_phase, x, Const(ε̇II_eff), Const(divV), Const(params))
        x = x .- J.derivs[1]\J.val
        if norm(J.val)<1e-10
            break
        end
    end

    # Recompute tensorial components using invariant scaling
    τ = ε̇_eff .* x[1]./ε̇II_eff
    return [τ[1], τ[2], τ[3], x[2]], x[3]
end

function single_phase_return_mapping()

    # Kinematics
    ε̇    = [0.1, -0.1, 0]
    divV = -0.05   

    # Parameters
    params = (
        G     = 1.0,
        K     = 3.0,
        C     = 1.0,
    )  
    
    # Invariants
    ε̇     = [ε̇[1], ε̇[2], ε̇[3], divV]

    # Consistent tangent operator ≡ ∂σ∂ϵ̇ 
    J = Enzyme.jacobian(Enzyme.ForwardWithPrimal, StressVector, ε̇, Const(params))
    display(J.derivs[1])
end

single_phase_return_mapping()

[tduretz]

It seems that the problem can be alleviated by using StaticArrays

using Enzyme, LinearAlgebra, StaticArrays

invII(x) = sqrt(1/2*x[1]^2 + 1/2*x[2]^2 + x[3]^2) 

# Function to be minimized via inner Newton iteration 
function residual_single_phase(x, ε̇II_eff, divV, P0, p)
    G, K, C, ϕ, ψ, ηvp, Δt = p.G, p.K, p.C, p.ϕ, p.ψ, p.ηvp, p.Δt
    eps   = -1e-13
    ηe    = G*Δt
    χe    = K*Δt
    τII, P, λ̇ = x[1], x[2], x[3]
    f      = τII  - C*cosd(ϕ) - P*sind(ϕ)
    return @SVector[ 
        ε̇II_eff  -  (τII)/2/ηe - λ̇*(f>=eps),
        divV     + (P - P0)/χe - λ̇*sind(ψ)*(f>=eps),
        (f - ηvp*λ̇)*(f>=eps) +  λ̇*1*(f<eps)
    ]
end

# Evaluate the stress vector involving the solution of a non-linear function
function StressVector(ϵ̇, τ0, P0, params)

    ε̇_eff, divV = ϵ̇[1:3],  ϵ̇[4]
    ε̇II_eff, τII = invII(ε̇_eff), invII(τ0)

    # Initialise inner solution array
    x = @SVector([τII, P0, 0.0])

    # Newton iteration
    for iter=1:10
        J = Enzyme.jacobian(Enzyme.ForwardWithPrimal, residual_single_phase, x, Const(ε̇II_eff), Const(divV), Const(P0), Const(params))
        x = x .- J.derivs[1]\J.val
        if norm(J.val)<1e-10
            break
        end
    end

    # Recompute tensorial components using invariant scaling
    τII, P, λ̇ = x[1], x[2], x[3]
    τ = ε̇_eff .* x[1]./ε̇II_eff
    return [τ[1], τ[2], τ[3], x[2]], x[3]
end

function single_phase_return_mapping()

    # Kinematics
    ε̇    = [0.1, -0.1, 0]
    divV = -0.05   

    # Initial conditions
    P    = 0.0
    τ    = [0.0, -0.0, 0]

    # Parameters
    params = (
        G     = 1.0,
        K     = 3.0,
        C     = 1.0,
        ϕ     = 35.0,
        ψ     = 35.0*0,
        ηvp   = 10.0*0,
        Δt    = 1.0,
    )  

    # Old guys
    P0 = P
    τ0 = τ
    
    # Invariants
    ε̇_eff = ε̇ + τ0/(2*params.G*params.Δt)
    ϵ̇     = [ε̇_eff[1], ε̇_eff[2], ε̇_eff[3], divV]

    # Consistent tangent operator ≡ ∂σ∂ϵ̇ 
    J = Enzyme.jacobian(Enzyme.ForwardWithPrimal, StressVector, ϵ̇, Const(τ0), Const(P0), Const(params))
    display(J.derivs[1])
end

single_phase_return_mapping()

[vchuravy]

So the crux seems to be that J = Enzyme.jacobian(Enzyme.ForwardWithPrimal, residual_single_phase, x, Const(ε̇II_eff), Const(divV), Const(params)) is not type-stable.

E.g. try @code_warntype StressVector(ε̇, params) this then trips up the second-order Enzyme.

This kinda bols down to Enzyme.onhot(zeros(3)) returning Tuple{Vararg{Vector{Float64}}}

[vchuravy]

So something like:

julia> function StressVector(ϵ̇, params)

           ε̇_eff, divV = ϵ̇[1:3], ϵ̇[4]
           ε̇II_eff = 1.0

           # Initialise inner solution array
           x = [1.0, 1.0, 0.0]
           shadows = ((
              [1.0, 0.0, 0.0],
              [0.0, 1.0, 0.0],
              [0.0, 0.0, 1.0]
           ),)

           # Newton iteration
           for iter=1:10
               J = Enzyme.jacobian(Enzyme.ForwardWithPrimal, residual_single_phase, x, Const(ε̇II_eff), Const(divV), Const(params); shadows)
               x .=- J.derivs[1]\J.val
               if norm(J.val)<1e-10
                   break
               end
           end

           # Recompute tensorial components using invariant scaling
           τ = ε̇_eff .* x[1]./ε̇II_eff
           return [τ[1], τ[2], τ[3], x[2]], x[3]
       end

Is better since it doesn't allocate shadows on each iteration and has the "nice" benefit of solving the type-stability issue.

Doesn't fully solve the nested issue though.

[vchuravy]

@wsmoses I see no way of making onehot(x:: Array) type stable since it will always depend on the length of a runtime value.

Could we add a variant of BatchDuplicated that does not use a Tuple, but rather an Array and we insert runtime checks that the length matches what we need?

[tduretz]

thanks for the hints. Why is the version with StaticArrays working? Do StaticArrays avoid the type instability?

[vchuravy]

Yes, for static arrays onehot(x) is typestables since we statically know the length of the shadow tuple we need to generate.

[h-spiess]

Hey. In this issue I think I have a similar issue: #2542.

Using StaticArrays makes Enzyme work, however the results are wrong. I will update the issue, but maybe check your second order gradient against e.g. ForwardDiff.

[wsmoses]

you can set chunksize = 1, I think and it ought help.

currently the internals of enzyme need a fixed width to compile the corresponding vectorized version of each derivative op so batchduplicated in its current form is really needed.

at one point we should support variable width chunk sizes, but thats a much larger and pervasive change