Implicit Function Adjoint Rules for Enzyme

[m-bossart]

I am trying to use Enzyme to differentiate a function which builds and solves a NonlinearProblem. I believ the issue is with the enzyme extension which is why I'm opening the issue here.

Expected behavior
I expect the SciMLSensitivity overload for NonlinearProblem to be called by Enzyme when differentiating.

Minimal Reproducible Example 👇

using SciMLSensitivity
using NonlinearSolve
using Enzyme
using Zygote
using NonlinearSolve
using Test

function nonlinearf(du, u, p)
    du[1] = u[1] .- p[1]
    nothing 
end 
p = [1.0]
function f(p)
    prob_nl = NonlinearProblem{true}(nonlinearf, [0.0], p)
    sol_nl = solve(prob_nl, TrustRegion(); abstol = 1e-9, reltol=1e-9)
    return sol_nl.u[1] * 2 
end 

dp = zeros(1)
dp_zygote = Zygote.gradient(f, p)[1]    #Zygote Works
Enzyme.autodiff(Reverse, f, Active, Duplicated(p, dp))  #Enzyme Fails

Error & Stacktrace ⚠️
I believe this is the problematic line: https://github.com/SciML/DiffEqBase.jl/blob/1658ed0c18e5e8e6776e7da67c52a56af5076fa7/ext/DiffEqBaseEnzymeExt.jl#L33

ERROR: MethodError: no method matching copyto!(::Float64, ::Base.Broadcast.Broadcasted{…})

Closest candidates are:
  copyto!(::AbstractArray, ::Base.Broadcast.Broadcasted{<:Base.Broadcast.AbstractArrayStyle{0}})
   @ Base broadcast.jl:959
  copyto!(::AbstractArray, ::Base.Broadcast.Broadcasted)
   @ Base broadcast.jl:956
  copyto!(::DiffEqArrayOperator, ::Any)
   @ SciMLBase C:\Users\Matt Bossart\.julia\packages\SciMLBase\SDjaO\src\operators\basic_operators.jl:127
  ...

Environment (please complete the following information):

• Output of using Pkg; Pkg.status()

  [d360d2e6] ChainRulesCore v1.23.0
  [b0b7db55] ComponentArrays v0.15.13
  [2b5f629d] DiffEqBase v6.151.1
  [7da242da] Enzyme v0.12.7
  [8913a72c] NonlinearSolve v3.12.0
  [7f7a1694] Optimization v3.25.1
  [42dfb2eb] OptimizationOptimisers v0.2.1
  [1dea7af3] OrdinaryDiffEq v6.78.0
  [f0f68f2c] PlotlyJS v0.18.13
  [295af30f] Revise v3.5.14
  [1ed8b502] SciMLSensitivity v7.59.0
  [e88e6eb3] Zygote v0.6.70

• Output of using Pkg; Pkg.status(; mode = PKGMODE_MANIFEST)

[47edcb42] ADTypes v1.2.1
  [621f4979] AbstractFFTs v1.5.0
  [1520ce14] AbstractTrees v0.4.5
  [7d9f7c33] Accessors v0.1.36
  [79e6a3ab] Adapt v4.0.4
  [66dad0bd] AliasTables v1.1.3
  [ec485272] ArnoldiMethod v0.4.0
  [4fba245c] ArrayInterface v7.10.0
  [4c555306] ArrayLayouts v1.9.3
  [bf4720bc] AssetRegistry v0.1.0
  [a9b6321e] Atomix v0.1.0
  [d1d4a3ce] BitFlags v0.1.8
  [62783981] BitTwiddlingConvenienceFunctions v0.1.5
  [ad839575] Blink v0.12.9
  [fa961155] CEnum v0.5.0
  [2a0fbf3d] CPUSummary v0.2.5
  [49dc2e85] Calculus v0.5.1
  [7057c7e9] Cassette v0.3.13
  [082447d4] ChainRules v1.66.0
  [d360d2e6] ChainRulesCore v1.23.0
  [fb6a15b2] CloseOpenIntervals v0.1.12
  [da1fd8a2] CodeTracking v1.3.5
  [944b1d66] CodecZlib v0.7.4
  [35d6a980] ColorSchemes v3.25.0
  [3da002f7] ColorTypes v0.11.5
  [c3611d14] ColorVectorSpace v0.10.0
  [5ae59095] Colors v0.12.11
  [38540f10] CommonSolve v0.2.4
  [bbf7d656] CommonSubexpressions v0.3.0
  [34da2185] Compat v4.15.0
  [b0b7db55] ComponentArrays v0.15.13
  [a33af91c] CompositionsBase v0.1.2
  [2569d6c7] ConcreteStructs v0.2.3
  [f0e56b4a] ConcurrentUtilities v2.4.1
  [88cd18e8] ConsoleProgressMonitor v0.1.2
  [187b0558] ConstructionBase v1.5.5
  [adafc99b] CpuId v0.3.1
  [9a962f9c] DataAPI v1.16.0
  [864edb3b] DataStructures v0.18.20
  [e2d170a0] DataValueInterfaces v1.0.0
  [8bb1440f] DelimitedFiles v1.9.1
  [2b5f629d] DiffEqBase v6.151.1
  [459566f4] DiffEqCallbacks v3.6.2
  [77a26b50] DiffEqNoiseProcess v5.21.0
  [163ba53b] DiffResults v1.1.0
  [b552c78f] DiffRules v1.15.1
  [b4f34e82] Distances v0.10.11
  [31c24e10] Distributions v0.25.108
  [ffbed154] DocStringExtensions v0.9.3
  [fa6b7ba4] DualNumbers v0.6.8
  [da5c29d0] EllipsisNotation v1.8.0
  [4e289a0a] EnumX v1.0.4
  [7da242da] Enzyme v0.12.7
  [f151be2c] EnzymeCore v0.7.2
  [460bff9d] ExceptionUnwrapping v0.1.10
  [d4d017d3] ExponentialUtilities v1.26.1
  [e2ba6199] ExprTools v0.1.10
  [7034ab61] FastBroadcast v0.2.8
  [9aa1b823] FastClosures v0.3.2
⌃ [29a986be] FastLapackInterface v2.0.3
  [1a297f60] FillArrays v1.11.0
  [6a86dc24] FiniteDiff v2.23.1
  [53c48c17] FixedPointNumbers v0.8.5
  [f6369f11] ForwardDiff v0.10.36
  [f62d2435] FunctionProperties v0.1.2
  [069b7b12] FunctionWrappers v1.1.3
  [77dc65aa] FunctionWrappersWrappers v0.1.3
  [de31a74c] FunctionalCollections v0.5.0
  [d9f16b24] Functors v0.4.10
  [0c68f7d7] GPUArrays v10.1.1
  [46192b85] GPUArraysCore v0.1.6
  [61eb1bfa] GPUCompiler v0.26.4
  [c145ed77] GenericSchur v0.5.4
  [86223c79] Graphs v1.11.0
  [cd3eb016] HTTP v1.10.8
  [9fb69e20] Hiccup v0.2.2
  [3e5b6fbb] HostCPUFeatures v0.1.16
  [34004b35] HypergeometricFunctions v0.3.23
  [7869d1d1] IRTools v0.4.14
  [615f187c] IfElse v0.1.1
  [d25df0c9] Inflate v0.1.4
  [3587e190] InverseFunctions v0.1.14
  [92d709cd] IrrationalConstants v0.2.2
  [82899510] IteratorInterfaceExtensions v1.0.0
  [692b3bcd] JLLWrappers v1.5.0
  [97c1335a] JSExpr v0.5.4
  [682c06a0] JSON v0.21.4
  [aa1ae85d] JuliaInterpreter v0.9.31
  [ccbc3e58] JumpProcesses v9.11.1
  [ef3ab10e] KLU v0.6.0
  [63c18a36] KernelAbstractions v0.9.19
  [ba0b0d4f] Krylov v0.9.6
  [5be7bae1] LBFGSB v0.4.1
⌅ [929cbde3] LLVM v6.6.3
  [b964fa9f] LaTeXStrings v1.3.1
  [10f19ff3] LayoutPointers v0.1.15
  [50d2b5c4] Lazy v0.15.1
⌅ [5078a376] LazyArrays v1.10.0
  [1d6d02ad] LeftChildRightSiblingTrees v0.2.0
  [2d8b4e74] LevyArea v1.0.0
  [d3d80556] LineSearches v7.2.0
  [7ed4a6bd] LinearSolve v2.30.0
  [2ab3a3ac] LogExpFunctions v0.3.27
  [e6f89c97] LoggingExtras v1.0.3
  [bdcacae8] LoopVectorization v0.12.170
  [6f1432cf] LoweredCodeUtils v2.4.6
  [1914dd2f] MacroTools v0.5.13
  [d125e4d3] ManualMemory v0.1.8
⌅ [a3b82374] MatrixFactorizations v2.2.0
  [bb5d69b7] MaybeInplace v0.1.2
  [739be429] MbedTLS v1.1.9
  [e1d29d7a] Missings v1.2.0
  [46d2c3a1] MuladdMacro v0.2.4
  [ffc61752] Mustache v1.0.19
  [a975b10e] Mux v1.0.2
  [d41bc354] NLSolversBase v7.8.3
  [2774e3e8] NLsolve v4.5.1
  [872c559c] NNlib v0.9.17
  [77ba4419] NaNMath v1.0.2
  [8913a72c] NonlinearSolve v3.12.0
  [d8793406] ObjectFile v0.4.1
  [510215fc] Observables v0.5.5
  [6fe1bfb0] OffsetArrays v1.14.0
  [4d8831e6] OpenSSL v1.4.3
  [429524aa] Optim v1.9.4
  [3bd65402] Optimisers v0.3.3
  [7f7a1694] Optimization v3.25.1
⌃ [bca83a33] OptimizationBase v1.0.0
  [42dfb2eb] OptimizationOptimisers v0.2.1
  [bac558e1] OrderedCollections v1.6.3
  [1dea7af3] OrdinaryDiffEq v6.78.0
  [90014a1f] PDMats v0.11.31
  [65ce6f38] PackageExtensionCompat v1.0.2
  [d96e819e] Parameters v0.12.3
  [69de0a69] Parsers v2.8.1
  [fa939f87] Pidfile v1.3.0
  [a03496cd] PlotlyBase v0.8.19
  [f0f68f2c] PlotlyJS v0.18.13
  [f2990250] PlotlyKaleido v2.2.4
  [e409e4f3] PoissonRandom v0.4.4
  [f517fe37] Polyester v0.7.14
  [1d0040c9] PolyesterWeave v0.2.1
  [85a6dd25] PositiveFactorizations v0.2.4
  [d236fae5] PreallocationTools v0.4.21
  [aea7be01] PrecompileTools v1.2.1
  [21216c6a] Preferences v1.4.3
  [33c8b6b6] ProgressLogging v0.1.4
  [92933f4c] ProgressMeter v1.10.0
  [43287f4e] PtrArrays v1.1.0
  [1fd47b50] QuadGK v2.9.4
  [74087812] Random123 v1.7.0
  [e6cf234a] RandomNumbers v1.5.3
  [c1ae055f] RealDot v0.1.0
  [3cdcf5f2] RecipesBase v1.3.4
  [731186ca] RecursiveArrayTools v3.19.0
  [f2c3362d] RecursiveFactorization v0.2.23
  [189a3867] Reexport v1.2.2
  [ae029012] Requires v1.3.0
  [ae5879a3] ResettableStacks v1.1.1
  [37e2e3b7] ReverseDiff v1.15.3
  [295af30f] Revise v3.5.14
  [79098fc4] Rmath v0.7.1
  [7e49a35a] RuntimeGeneratedFunctions v0.5.13
  [94e857df] SIMDTypes v0.1.0
  [476501e8] SLEEFPirates v0.6.42
  [0bca4576] SciMLBase v2.38.0
  [c0aeaf25] SciMLOperators v0.3.8
  [1ed8b502] SciMLSensitivity v7.59.0
  [53ae85a6] SciMLStructures v1.2.0
  [6c6a2e73] Scratch v1.2.1
  [efcf1570] Setfield v1.1.1
  [777ac1f9] SimpleBufferStream v1.1.0
  [727e6d20] SimpleNonlinearSolve v1.8.0
  [699a6c99] SimpleTraits v0.9.4
  [ce78b400] SimpleUnPack v1.1.0
  [a2af1166] SortingAlgorithms v1.2.1
  [47a9eef4] SparseDiffTools v2.19.0
  [dc90abb0] SparseInverseSubset v0.1.2
  [e56a9233] Sparspak v0.3.9
  [276daf66] SpecialFunctions v2.4.0
  [aedffcd0] Static v0.8.10
  [0d7ed370] StaticArrayInterface v1.5.0
  [90137ffa] StaticArrays v1.9.4
  [1e83bf80] StaticArraysCore v1.4.2
  [82ae8749] StatsAPI v1.7.0
  [2913bbd2] StatsBase v0.34.3
  [4c63d2b9] StatsFuns v1.3.1
  [789caeaf] StochasticDiffEq v6.65.1
  [7792a7ef] StrideArraysCore v0.5.6
  [09ab397b] StructArrays v0.6.18
  [53d494c1] StructIO v0.3.0
  [2efcf032] SymbolicIndexingInterface v0.3.21
  [3783bdb8] TableTraits v1.0.1
  [bd369af6] Tables v1.11.1
  [62fd8b95] TensorCore v0.1.1
  [5d786b92] TerminalLoggers v0.1.7
  [8290d209] ThreadingUtilities v0.5.2
  [a759f4b9] TimerOutputs v0.5.24
  [9f7883ad] Tracker v0.2.34
  [3bb67fe8] TranscodingStreams v0.10.8
  [d5829a12] TriangularSolve v0.2.0
  [410a4b4d] Tricks v0.1.8
  [781d530d] TruncatedStacktraces v1.4.0
  [5c2747f8] URIs v1.5.1
  [3a884ed6] UnPack v1.0.2
  [013be700] UnsafeAtomics v0.2.1
  [d80eeb9a] UnsafeAtomicsLLVM v0.1.4
  [3d5dd08c] VectorizationBase v0.21.67
  [19fa3120] VertexSafeGraphs v0.2.0
  [0f1e0344] WebIO v0.8.21
  [104b5d7c] WebSockets v1.6.0
  [cc8bc4a8] Widgets v0.6.6
  [e88e6eb3] Zygote v0.6.70
  [700de1a5] ZygoteRules v0.2.5
⌅ [7cc45869] Enzyme_jll v0.0.111+0
  [1d5cc7b8] IntelOpenMP_jll v2024.1.0+0
  [f7e6163d] Kaleido_jll v0.2.1+0
  [dad2f222] LLVMExtra_jll v0.0.29+0
  [81d17ec3] L_BFGS_B_jll v3.0.1+0
  [856f044c] MKL_jll v2024.1.0+0
  [458c3c95] OpenSSL_jll v3.0.13+1
  [efe28fd5] OpenSpecFun_jll v0.5.5+0
  [f50d1b31] Rmath_jll v0.4.2+0
  [1317d2d5] oneTBB_jll v2021.12.0+0
  [0dad84c5] ArgTools v1.1.1
  [56f22d72] Artifacts
  [2a0f44e3] Base64
  [ade2ca70] Dates
  [8ba89e20] Distributed
  [f43a241f] Downloads v1.6.0
  [7b1f6079] FileWatching
  [9fa8497b] Future
  [b77e0a4c] InteractiveUtils
  [4af54fe1] LazyArtifacts
  [b27032c2] LibCURL v0.6.4
  [76f85450] LibGit2
  [8f399da3] Libdl
  [37e2e46d] LinearAlgebra
  [56ddb016] Logging
  [d6f4376e] Markdown
  [a63ad114] Mmap
  [ca575930] NetworkOptions v1.2.0
  [44cfe95a] Pkg v1.10.0
  [de0858da] Printf
  [3fa0cd96] REPL
  [9a3f8284] Random
  [ea8e919c] SHA v0.7.0
  [9e88b42a] Serialization
  [1a1011a3] SharedArrays
  [6462fe0b] Sockets
  [2f01184e] SparseArrays v1.10.0
  [10745b16] Statistics v1.10.0
  [4607b0f0] SuiteSparse
  [fa267f1f] TOML v1.0.3
  [a4e569a6] Tar v1.10.0
  [8dfed614] Test
  [cf7118a7] UUIDs
  [4ec0a83e] Unicode
  [e66e0078] CompilerSupportLibraries_jll v1.1.0+0
  [deac9b47] LibCURL_jll v8.4.0+0
  [e37daf67] LibGit2_jll v1.6.4+0
  [29816b5a] LibSSH2_jll v1.11.0+1
  [c8ffd9c3] MbedTLS_jll v2.28.2+1
  [14a3606d] MozillaCACerts_jll v2023.1.10
  [4536629a] OpenBLAS_jll v0.3.23+4
  [05823500] OpenLibm_jll v0.8.1+2
  [bea87d4a] SuiteSparse_jll v7.2.1+1
  [83775a58] Zlib_jll v1.2.13+1
  [8e850b90] libblastrampoline_jll v5.8.0+1
  [8e850ede] nghttp2_jll v1.52.0+1
  [3f19e933] p7zip_jll v17.4.0+2

• Output of versioninfo()

Julia Version 1.10.2
Commit bd47eca2c8 (2024-03-01 10:14 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Windows (x86_64-w64-mingw32)
  CPU: 16 × 11th Gen Intel(R) Core(TM) i7-11800H @ 2.30GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-15.0.7 (ORCJIT, tigerlake)
Threads: 1 default, 0 interactive, 1 GC (on 16 virtual cores)
Environment:
  JULIA_EDITOR = code
  JULIA_NUM_THREADS = 

[ChrisRackauckas]

Note that NonlinearSolve is currently missing Enzyme overloads. Only LinearSolve and ODEProblems have gotten that treatment so far.

[wsmoses]

@ChrisRackauckas I can help with the EnzymeRules infra, is there an adjoint function lying around for this. Similarly what would be the right function to hook into with a rule

[ChrisRackauckas]

It's just like the ODE infrastructure in that it just lowers to call SciMLSensitivity for all of the real work. The current dispatch is in DiffEqBase.jl right next to the ODE one.

[wsmoses]

Sorry you're going to have to give me some more concrete pointers to functions which should have a rule.

[m-bossart]

I believe this is the dispatch Chris is referring to:
(https://github.com/SciML/DiffEqBase.jl/blob/08fbaf0561d797359027849bac7797ce10948846/src/solve.jl#L1066C1-L1083C4)

For the ODE case, the enzyme rule is for the solve_up function, so I'm guessing the same should be true for the case of nonlinear problems:
https://github.com/SciML/DiffEqBase.jl/blob/08fbaf0561d797359027849bac7797ce10948846/ext/DiffEqBaseEnzymeExt.jl#L9C1-L37C4

I'm not sure though, @ChrisRackauckas can you confirm?

[ChrisRackauckas]

Yup and then the adjoint just passes it on: https://github.com/SciML/DiffEqBase.jl/blob/08fbaf0561d797359027849bac7797ce10948846/ext/DiffEqBaseChainRulesCoreExt.jl#L22C28-L29 . Looking at the ODE one, the NonlinearProblem one might just work if NonlinearProblem is made a mutable struct.

[m-bossart]

I tried making NonlinearProblem a mutable struct but am still getting an error:

Closest candidates are:
  copyto!(::OrdinaryDiffEq.ArrayFuse, ::Base.Broadcast.Broadcasted)
   @ OrdinaryDiffEq C:\Users\Matt Bossart\.julia\packages\OrdinaryDiffEq\tAI61\src\wrappers.jl:26
  copyto!(::AbstractArray, ::Base.Broadcast.Broadcasted{<:Base.Broadcast.AbstractArrayStyle{0}})
   @ Base broadcast.jl:959
  copyto!(::AbstractArray, ::Base.Broadcast.Broadcasted)
   @ Base broadcast.jl:956
  ...

It is failing at this line:
https://github.com/SciML/DiffEqBase.jl/blob/1658ed0c18e5e8e6776e7da67c52a56af5076fa7/ext/DiffEqBaseEnzymeExt.jl#L33

[wsmoses]

What is the type of the variable being zero's there

Enzyme semantics require it be zeroed but it looks like it may not be defined for that type?

[m-bossart]

The type of the variable being zero'd is Float64. The error I hit above as the same error you get if you do (trying to broadcast a value to a float):

x = 1.0
x .= 0.0

The part I'm unclear on is why this works for an ODEProblem and not a NonlinearProblem. @ChrisRackauckas is it expected that the return from _solve_adjoint would be different for those two?

[wsmoses]

Oh that's definitely it.

Here's a solution. Replace the deepcopy with an Enzyme.make_zero and get rid of the for loop