Prep for Enzyme breaking Change

[wsmoses]

Enzyme has a pending breaking change. Among other things, we're trying to improve sugar of gradient/jacobian/etc.

I'm curious if you can take a look and see if there's some functionality we can add that would simplify your use here.

[Vaibhavdixit02]

I don't use the gradient/jacobian/etc functions here anymore, autodiff actually has some features that are nice to work with for me here (mentioned below). But there's two main limitations of our implementation here https://github.com/SciML/OptimizationBase.jl/blob/main/ext/OptimizationEnzymeExt.jl

1. I have tried to roll in mini-batching and constant parameters in the v2 of OptimizationBase into the second argument to the loss/grad/etc, though the Enzyme interface for it is pretty spot on with autodiff and I expected to just work, unfortunately it doesn't, I am not sure I completely understand why. The MWE (with OptimizationBase) is

using MLUtils, OptimizationBase

x = rand(10000)
    y = sin.(x)
    data = MLUtils.DataLoader((x, y), batchsize = 100)

    function loss(coeffs, data)
        ypred = [evalpoly(data[1][i], coeffs) for i in eachindex(data[1])]
        return sum(abs2, ypred .- data[2])
    end
optf = OptimizationFunction(loss, AutoEnzyme())
    optf = OptimizationBase.instantiate_function(
        optf, rand(3), AutoEnzyme(), iterate(data)[1], g = true, fg = true)
    G0 = zeros(3)
    @test_broken optf.grad(G0, ones(3))

ERROR: Constant memory is stored (or returned) to a differentiable variable.
As a result, Enzyme cannot provably ensure correctness and throws this error.
This might be due to the use of a constant variable as temporary storage for active memory (https://enzyme.mit.edu/julia/stable/faq/#Activity-of-temporary-storage).
If Enzyme should be able to prove this use non-differentable, open an issue!
To work around this issue, either:
 a) rewrite this variable to not be conditionally active (fastest, but requires a code change), or
 b) set Enzyme.API.runtimeActivity!(true) immediately after loading Enzyme (which maintains correctness, but may slightly reduce performance).
Mismatched activity for:   store {} addrspace(10)* %unbox.unpack, {} addrspace(10)** %.fca.0.1.0.gep, align 8, !dbg !26, !noalias !56 const val:   %unbox.unpack = load {} addrspace(10)*, {} addrspace(10)* addrspace(11)* %unbox.elt, align 8, !dbg !26, !tbaa !23, !alias.scope !27, !noalias !30
 value=Unknown object of type Vector{Float64}
 llvalue=  %unbox.unpack = load {} addrspace(10)*, {} addrspace(10)* addrspace(11)* %unbox.elt, align 8, !dbg !26, !tbaa !23, !alias.scope !27, !noalias !30

Stacktrace:
 [1] loss
   @ ~/OptimizationBase.jl/test/adtests.jl:1105

Stacktrace:
  [1] loss
    @ ~/OptimizationBase.jl/test/adtests.jl:1105
  [2] firstapply
    @ ~/OptimizationBase.jl/ext/OptimizationEnzymeExt.jl:12 [inlined]
  [3] firstapply
    @ ~/OptimizationBase.jl/ext/OptimizationEnzymeExt.jl:0 [inlined]
  [4] diffejulia_firstapply_6167_inner_1wrap
    @ ~/OptimizationBase.jl/ext/OptimizationEnzymeExt.jl:0
  [5] macro expansion
    @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:7187 [inlined]
  [6] enzyme_call
    @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:6794 [inlined]
  [7] CombinedAdjointThunk
    @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:6671 [inlined]
  [8] autodiff
    @ ~/.julia/packages/Enzyme/TiboG/src/Enzyme.jl:320 [inlined]
  [9] grad
    @ ~/OptimizationBase.jl/ext/OptimizationEnzymeExt.jl:99 [inlined]
 [10] (::OptimizationEnzymeExt.var"#grad#24"{Tuple{…}, OptimizationFunction{…}})(res::Vector{Float64}, θ::Vector{Float64})
    @ OptimizationEnzymeExt ~/OptimizationBase.jl/ext/OptimizationEnzymeExt.jl:98
 [11] top-level scope
    @ ~/OptimizationBase.jl/test/adtests.jl:1162

2. The lack of clear documentation on using sparsity patterns has prevented me from being able to implement it here, DI promises some solutions but they don't have a way to handle constants well yet. This makes it much slower than sparsity aware FD and RD so it would be nice to have that working, here's a recent plot for reference on this https://docs.sciml.ai/SciMLBenchmarksOutput/dev/OptimizationFrameworks/clnlbeam/ benchmark - this may be something that the interface changes you mentioned could make a bit better?

[Vaibhavdixit02]

Fixed by #110