Computing third order derivatives

[Arpit-Babbar]

I am able to go up to second order derivatives with Enzyme. However, when I use the same approach to go for the third derivative, I get an error. Here is a minimal working example to show that. The last function call (supposed to compute the third order derivative) gives the error.

using Enzyme
using StaticArrays
function g(x) # Some function to be differentiated
    x[1]^4 + x[1]^7 + x[2]^4 + x[1]^2 * x[2]^2
end

# Enzyme.jl derivatives
dg(x, dx) = autodiff_deferred(Forward, g, Duplicated, Duplicated(x, dx))
ddg(x, dx, ddx) = autodiff_deferred(Forward, dg, Duplicated, Duplicated(x, dx), Duplicated(dx, ddx))
dddg(x, dx, ddx, dddx) = autodiff_deferred(Forward, ddg, Duplicated, Duplicated(x, dx),
                                           Duplicated(dx, ddx), Duplicated(ddx, dddx))
                                           
# Get random inputs for x and some directions for the directional derivatives
x = SVector{2}(rand(2)); dx = SVector{2}(rand(2)); ddx = SVector{2}(rand(2)); dddx = SVector{2}(rand(2))

# Call the functions with inputs
dg(x, dx) # Works
ddg(x, dx, ddx) # Works
dddg(x, dx, ddx, dddx) # Error

The error message is below.

Error message

ERROR: Enzyme execution failed.
Enzyme compilation failed.

No forward mode derivative found for deferred_codegen
 at context:   %5 = call i64 @deferred_codegen(i64 noundef 139926815969256) #5, !dbg !19, !noalias !8

Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/Enzyme/fS5Q3/src/compiler.jl:7113
 [2] deferred_codegen (repeats 3 times)
   @ ~/.julia/packages/Enzyme/fS5Q3/src/compiler.jl:7072
 [3] autodiff_deferred
   @ ~/.julia/packages/Enzyme/fS5Q3/src/Enzyme.jl:524
 [4] autodiff_deferred
   @ ~/.julia/packages/Enzyme/fS5Q3/src/Enzyme.jl:538
 [5] dg
   @ ~/repositories/2024_ad_lw_dev/TenkaiAD/misc/third_derivative_test_arpit.jl:78
 [6] dg
   @ ~/repositories/2024_ad_lw_dev/TenkaiAD/misc/third_derivative_test_arpit.jl:0


Stacktrace:
  [1] throwerr(cstr::Cstring)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/fS5Q3/src/compiler.jl:1797
  [2] macro expansion
    @ ~/.julia/packages/Enzyme/fS5Q3/src/compiler.jl:7113 [inlined]
  [3] deferred_codegen (repeats 3 times)
    @ ~/.julia/packages/Enzyme/fS5Q3/src/compiler.jl:7072 [inlined]
  [4] autodiff_deferred
    @ ~/.julia/packages/Enzyme/fS5Q3/src/Enzyme.jl:524 [inlined]
  [5] autodiff_deferred
    @ ~/.julia/packages/Enzyme/fS5Q3/src/Enzyme.jl:538 [inlined]
  [6] dg
    @ ~/repositories/2024_ad_lw_dev/TenkaiAD/misc/third_derivative_test_arpit.jl:78 [inlined]
  [7] dg
    @ ~/repositories/2024_ad_lw_dev/TenkaiAD/misc/third_derivative_test_arpit.jl:0 [inlined]
  [8] fwddiffejulia_dg_5482_inner_1wrap
    @ ~/repositories/2024_ad_lw_dev/TenkaiAD/misc/third_derivative_test_arpit.jl:0 [inlined]
  [9] macro expansion
    @ ~/.julia/packages/Enzyme/fS5Q3/src/compiler.jl:6837 [inlined]
 [10] enzyme_call
    @ ~/.julia/packages/Enzyme/fS5Q3/src/compiler.jl:6437 [inlined]
 [11] ForwardModeThunk
    @ ~/.julia/packages/Enzyme/fS5Q3/src/compiler.jl:6317 [inlined]
 [12] autodiff_deferred
    @ ~/.julia/packages/Enzyme/fS5Q3/src/Enzyme.jl:526 [inlined]
 [13] autodiff_deferred
    @ ~/.julia/packages/Enzyme/fS5Q3/src/Enzyme.jl:538 [inlined]
 [14] ddg
    @ ~/repositories/2024_ad_lw_dev/TenkaiAD/misc/third_derivative_test_arpit.jl:79 [inlined]
 [15] ddg
    @ ~/repositories/2024_ad_lw_dev/TenkaiAD/misc/third_derivative_test_arpit.jl:0 [inlined]
 [16] fwddiffejulia_ddg_5479_inner_1wrap
    @ ~/repositories/2024_ad_lw_dev/TenkaiAD/misc/third_derivative_test_arpit.jl:0
 [17] macro expansion
    @ ~/.julia/packages/Enzyme/fS5Q3/src/compiler.jl:6837 [inlined]
 [18] enzyme_call
    @ ~/.julia/packages/Enzyme/fS5Q3/src/compiler.jl:6437 [inlined]
 [19] ForwardModeThunk
    @ ~/.julia/packages/Enzyme/fS5Q3/src/compiler.jl:6317 [inlined]
 [20] autodiff_deferred
    @ ~/.julia/packages/Enzyme/fS5Q3/src/Enzyme.jl:526 [inlined]
 [21] autodiff_deferred
    @ ~/.julia/packages/Enzyme/fS5Q3/src/Enzyme.jl:538 [inlined]
 [22] dddg(x::SVector{2, Float64}, dx::SVector{2, Float64}, ddx::SVector{2, Float64}, dddx::SVector{2, Float64})
    @ Main ~/repositories/2024_ad_lw_dev/TenkaiAD/misc/third_derivative_test_arpit.jl:80
 [23] top-level scope
    @ ~/repositories/2024_ad_lw_dev/TenkaiAD/misc/third_derivative_test_arpit.jl:85

I am using Enzyme v0.12.27.

versioninfo()

Julia Version 1.10.3
Commit 0b4590a5507 (2024-04-30 10:59 UTC)
Build Info:
Official https://julialang.org/ release
Platform Info:
OS: Linux (x86_64-linux-gnu)
CPU: 80 × Intel(R) Xeon(R) Gold 6148 CPU @ 2.40GHz
WORD_SIZE: 64
LIBM: libopenlibm
LLVM: libLLVM-15.0.7 (ORCJIT, skylake-avx512)
Threads: 1 default, 0 interactive, 1 GC (on 80 virtual cores)
Environment:
JULIA_EDITOR =
JULIA_NUM_THREADS = 1

[vchuravy]

This is likely the same as #1173, the goal is that JuliaGPU/GPUCompiler.jl#599 will fix this.

[wsmoses]

Closing in favor of JuliaGPU/GPUCompiler.jl#629

[wsmoses]

This will now work nicely automatically once #2161 lands!

using Enzyme
using StaticArrays
function g(x) # Some function to be differentiated
    x[1]^4 + x[1]^7 + x[2]^4 + x[1]^2 * x[2]^2
end

# Enzyme.jl derivatives
dg(x, dx) = autodiff(Forward, g, Duplicated, Duplicated(x, dx))
ddg(x, dx, ddx) = autodiff(Forward, dg, Duplicated, Duplicated(x, dx), Duplicated(dx, ddx))
dddg(x, dx, ddx, dddx) = autodiff(Forward, ddg, Duplicated, Duplicated(x, dx),
                                           Duplicated(dx, ddx), Duplicated(ddx, dddx))
                                           
# Get random inputs for x and some directions for the directional derivatives
x = SVector{2}(rand(2)); dx = SVector{2}(rand(2)); ddx = SVector{2}(rand(2)); dddx = SVector{2}(rand(2))

# Call the functions with inputs
dg(x, dx) # Works
ddg(x, dx, ddx) # Works
dddg(x, dx, ddx, dddx) # Works now

[wsmoses]

using Enzyme

function d(f)
    return x -> (autodiff(Forward, Const(f), Duplicated(x, 1.0))[1])
end

@show sin(2.0)
@show d(sin)(2.0)
@show d(d(sin))(2.0)
@show d(d(d(sin)))(2.0)
@show d(d(d(d(sin))))(2.0)
@show d(d(d(d(d(sin)))))(2.0)
@show d(d(d(d(d(d(sin))))))(2.0)

[vchuravy]

Finally fixed by #2161

[Arpit-Babbar]

Thanks for all your effort. I tested it for my application involving third derivatives, and it works exactly as I would like it to!