Automate use of deferred in Higher order derivatives

[AtilaSaraiva]

Hello folks,

I tried to make calculate a second order derivative using Enzyme and wasn't able to. Is it possible?

Here is the code

using Enzyme


# Impure implementation of the f(x) = x^3 function
function f!(x::Vector{<:Number}, y::Vector{<:Number})
    @assert length(x) == length(y)
    for i=1:length(x)
        y[i] = x[i]^3
    end
    return nothing
end

x = Float64[1, 2, 3, 4] # input values
y = zeros(4)            # output values
dx = zeros(4)           # ∂f/∂x
dy = ones(4)            # dummy accumulator for ∂f/∂y

f!(x,y)

@show y - x.^3

autodiff(Reverse, f!, Const, Duplicated(x,dx), Duplicated(y, dy))

@show dx - 3*x.^2

function ∂f_∂x!(x::Vector{<:Number}, dx::Vector{<:Number})
    y = zeros(4)            # output values
    dy = ones(4)            # dummy accumulator for ∂f/∂y
    autodiff(Reverse, f!, Const, Duplicated(x,dx), Duplicated(y, dy))
    return nothing
end

dx = zeros(4)

∂f_∂x!(x, dx)

@show dx - 3*x.^2

dx2 = zeros(4)

# Trying to calculate a second order derivative with Enzyme
autodiff(Reverse, ∂f_∂x!, Const, Duplicated(x,dx2), Duplicated(y, dy))

@show dx2

When I try to run the second derivative the Julia REPL crashes and shows an allocation error. I can try to paste it here but it exceeds the maximum number of characters, so I'd have to upload a file.

[wsmoses]

It is! The inner derivative calls must use autodiff_deferred though

[MilesCranmer]

Just ran into this as well with the following error followed by a __pthread_kill which crashes my Julia

Assertion failed: (whatType(argType, Mode) == DIFFE_TYPE::DUP_ARG || whatType(argType, Mode) == DIFFE_TYPE::CONSTANT), function recursivelyHandleSubfunction, file /workspace/srcdir/Enzyme/enzyme/Enzyme/AdjointGenerator.h, line 7109.

Using autodiff_deferred fixed it, thanks!

It would really improve the user experience if this low-level error was caught and handled smoothly with a high-level error that points the user to autodiff_deferred, as I had to hunt around for this issue.

[MilesCranmer]

Also. Would it be possible in the future to avoid the need for calling autodiff_deferred manually? Or is it a fundamental design limitation?

I am just thinking this might cause issues in the distant future when Enzyme.jl is more commonplace, if someone differentiates a large program that has an internal Enzyme.autodiff call hiding somewhere. They would have to reach into the internals and replace those calls manually.

[wsmoses]

I actually slacked @aviatesk a bit earlier to ask if we could use AbstractInterpreter to do this automatically.

@aviatesk thoughts on this?

cc @vchuravy

[AtilaSaraiva]

It would be nice to have an option in autodiff_deferred to set the order of the derivative as well, like if I wanted a second or a third derivative with relation to a variable. Thank you guys for the amazing work you are doing.

[cgeoga]

Apologies if this would be better discussed in a new issue, but I'm also trying to play with higher-order derivatives and having some issues. My primary concern is non-allocating derivatives, and so I have the following little code setup:

module EnzymeHigherOrder

  using Enzyme
  
  struct DefDerivative{F}
    fn::F
  end
  function (dd::DefDerivative{F})(x) where{F} 
    autodiff_deferred(Forward, dd.fn, DuplicatedNoNeed, 
                      Duplicated(x, one(x)))[1]
  end

  struct Derivative{F}
    fn::F
  end
  function (dv::Derivative{F})(x) where{F} 
    autodiff(Forward, dv.fn, DuplicatedNoNeed, Duplicated(x, one(x)))[1]
  end

  derivative(fn::F, ::Val{0}) where{F}   = fn
  derivative(fn::F, ::Val{1}) where{F}   = Derivative(fn)
  derivative(fn::F, ::Val{N}) where{F,N} = derivative(DefDerivative(fn), Val(N-1))

end

And I can use this to get second derivatives as

const ddcos = EnzymeHigherOrder.derivative(cos, Val(2))
@btime ddcos($(1.1)); # no allocations, great perf, woohoo!

But when I ask for a third derivative I get an error:

const dddcos = EnzymeHigherOrder.derivative(cos, Val(3))
dddcos(1.1)

gives

ERROR: Enzyme execution failed.
Enzyme: Not yet implemented, mixed activity for jl_new_struct constants=Bool[1, 0]   %5 = call noalias nonnull {} 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 %0, {} addrspace(10)* nonnull %3) #11

Stacktrace:
 [1] throwerr(cstr::Cstring)
   @ Enzyme.Compiler ~/.julia/packages/Enzyme/rbuCz/src/compiler.jl:3044
 [2] macro expansion
   @ ~/.julia/packages/Enzyme/rbuCz/src/compiler.jl:9778 [inlined]
 [3] enzyme_call
   @ ~/.julia/packages/Enzyme/rbuCz/src/compiler.jl:9456 [inlined]
 [4] ForwardModeThunk
   @ ~/.julia/packages/Enzyme/rbuCz/src/compiler.jl:9422 [inlined]
 [5] autodiff
   @ ~/.julia/packages/Enzyme/rbuCz/src/Enzyme.jl:330 [inlined]
 [6] autodiff
   @ ~/.julia/packages/Enzyme/rbuCz/src/Enzyme.jl:222 [inlined]
 [7] (::Main.EnzymeHigherOrder.Derivative{Main.EnzymeHigherOrder.DefDerivative{Main.EnzymeHigherOrder.DefDerivative{typeof(cos)}}})(x::Float64)
   @ Main.EnzymeHigherOrder ~/Scratch/higherorderad/enzyme_ho.jl:18
 [8] top-level scope
   @ REPL[3]:1

Does this mean that autodiff_deferred can't differentiate other autodiff_deferred? So that if I want a third derivative in forward-mode I need to do something that looks like autodiff(autodiff(autodiff_deferred))?

Please let me know if this should be its own issue!

[wsmoses]

@cgeoga can you open a separate issue

[wsmoses]

Should be resolved by #1005

[vchuravy]

Actual PR was #1839 ?