AutoZygote incompatible with multiple shooting objective

Describe the bug 🐞

The cost function returned by multiple_shooting_objective is not differentiable by Zygote, which prevents using the AutoZygote AD setting in optimisation.

Expected behavior Zygote succesfully computes the gradient of the cost function. The optimisation finishes succesfully.

Minimal Reproducible Example 👇

using DifferentialEquations
using DiffEqParamEstim
using Optimization
using OptimizationOptimJL
using SciMLSensitivity
using Zygote

f(u, p, t) = p.*u

u0 = [1.0, 2.0]
p0 = [-1.0, -3.0]
tspan = (0.0, 10.0)
prob = ODEProblem(f, u0, tspan, p0)
realsol = solve(prob, saveat = tspan[1]:tspan[2])

pguess = rand(2) * 10 .- 10

cost_function = multiple_shooting_objective(
    prob, Tsit5(),
    L2Loss(tspan[1]:1.0:tspan[2], Array(realsol)),
    Optimization.AutoZygote(),
    maxiters=100)

optprob = Optimization.OptimizationProblem(cost_function, vcat(zeros(10), pguess))

Zygote.gradient(cost_function, vcat(zeros(10), pguess))

result_ode = Optimization.solve(optprob, BFGS())

Error & Stacktrace ⚠️ The stacktrace was too long to paste here. Error is shown below

ERROR: DimensionMismatch: variable with size(x) == (2, 3) cannot have a gradient with size(dx) == (3,)

Environment (please complete the following information):

Output of using Pkg; Pkg.status()

Project estimateRC v0.1.0
Status `~/tmp/estimateRC/Project.toml`
  [1130ab10] DiffEqParamEstim v2.1.0
  [0c46a032] DifferentialEquations v7.11.0
  [7f7a1694] Optimization v3.19.3
  [36348300] OptimizationOptimJL v0.1.13
  [0bca4576] SciMLBase v2.9.0
  [1ed8b502] SciMLSensitivity v7.47.0
  [e88e6eb3] Zygote v0.6.67

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

Project estimateRC v0.1.0
Status `~/tmp/estimateRC/Manifest.toml`
  [47edcb42] ADTypes v0.2.5
  [621f4979] AbstractFFTs v1.5.0
  [1520ce14] AbstractTrees v0.4.4
  [7d9f7c33] Accessors v0.1.33
  [79e6a3ab] Adapt v3.7.1
  [ec485272] ArnoldiMethod v0.2.0
  [4fba245c] ArrayInterface v7.6.1
  [4c555306] ArrayLayouts v1.4.3
  [a9b6321e] Atomix v0.1.0
  [aae01518] BandedMatrices v1.2.1
  [62783981] BitTwiddlingConvenienceFunctions v0.1.5
  [764a87c0] BoundaryValueDiffEq v5.4.0
⌅ [fa961155] CEnum v0.4.2
  [2a0fbf3d] CPUSummary v0.2.4
  [49dc2e85] Calculus v0.5.1
  [7057c7e9] Cassette v0.3.12
  [082447d4] ChainRules v1.58.0
  [d360d2e6] ChainRulesCore v1.18.0
  [fb6a15b2] CloseOpenIntervals v0.1.12
  [38540f10] CommonSolve v0.2.4
  [bbf7d656] CommonSubexpressions v0.3.0
  [34da2185] Compat v4.10.0
  [a33af91c] CompositionsBase v0.1.2
  [2569d6c7] ConcreteStructs v0.2.3
  [88cd18e8] ConsoleProgressMonitor v0.1.2
  [187b0558] ConstructionBase v1.5.4
  [adafc99b] CpuId v0.3.1
  [9a962f9c] DataAPI v1.15.0
  [864edb3b] DataStructures v0.18.15
  [e2d170a0] DataValueInterfaces v1.0.0
  [bcd4f6db] DelayDiffEq v5.43.1
  [8bb1440f] DelimitedFiles v1.9.1
  [39dd38d3] Dierckx v0.5.3
  [2b5f629d] DiffEqBase v6.141.0
  [459566f4] DiffEqCallbacks v2.34.0
  [77a26b50] DiffEqNoiseProcess v5.19.0
  [1130ab10] DiffEqParamEstim v2.1.0
  [163ba53b] DiffResults v1.1.0
  [b552c78f] DiffRules v1.15.1
  [0c46a032] DifferentialEquations v7.11.0
  [b4f34e82] Distances v0.10.10
  [31c24e10] Distributions v0.25.103
  [ffbed154] DocStringExtensions v0.9.3
  [fa6b7ba4] DualNumbers v0.6.8
  [da5c29d0] EllipsisNotation v1.8.0
  [4e289a0a] EnumX v1.0.4
  [7da242da] Enzyme v0.11.10
  [f151be2c] EnzymeCore v0.6.3
  [d4d017d3] ExponentialUtilities v1.25.0
  [e2ba6199] ExprTools v0.1.10
  [7034ab61] FastBroadcast v0.2.8
  [9aa1b823] FastClosures v0.3.2
  [29a986be] FastLapackInterface v2.0.0
  [1a297f60] FillArrays v1.7.0
  [6a86dc24] FiniteDiff v2.21.1
  [f6369f11] ForwardDiff v0.10.36
  [f62d2435] FunctionProperties v0.1.0
  [069b7b12] FunctionWrappers v1.1.3
  [77dc65aa] FunctionWrappersWrappers v0.1.3
  [d9f16b24] Functors v0.4.5
  [0c68f7d7] GPUArrays v9.1.0
  [46192b85] GPUArraysCore v0.1.5
  [61eb1bfa] GPUCompiler v0.25.0
  [c145ed77] GenericSchur v0.5.3
  [86223c79] Graphs v1.9.0
  [3e5b6fbb] HostCPUFeatures v0.1.16
  [34004b35] HypergeometricFunctions v0.3.23
  [7869d1d1] IRTools v0.4.11
  [615f187c] IfElse v0.1.1
  [d25df0c9] Inflate v0.1.4
  [18e54dd8] IntegerMathUtils v0.1.2
  [3587e190] InverseFunctions v0.1.12
  [92d709cd] IrrationalConstants v0.2.2
  [82899510] IteratorInterfaceExtensions v1.0.0
  [692b3bcd] JLLWrappers v1.5.0
  [ccbc3e58] JumpProcesses v9.8.0
  [ef3ab10e] KLU v0.4.1
  [63c18a36] KernelAbstractions v0.9.13
  [ba0b0d4f] Krylov v0.9.4
  [929cbde3] LLVM v6.4.0
  [73f95e8e] LatticeRules v0.0.1
  [10f19ff3] LayoutPointers v0.1.15
  [50d2b5c4] Lazy v0.15.1
  [1d6d02ad] LeftChildRightSiblingTrees v0.2.0
  [2d8b4e74] LevyArea v1.0.0
  [d3d80556] LineSearches v7.2.0
  [7ed4a6bd] LinearSolve v2.20.0
  [2ab3a3ac] LogExpFunctions v0.3.26
  [e6f89c97] LoggingExtras v1.0.3
  [bdcacae8] LoopVectorization v0.12.166
  [1914dd2f] MacroTools v0.5.11
  [d125e4d3] ManualMemory v0.1.8
  [e1d29d7a] Missings v1.1.0
  [46d2c3a1] MuladdMacro v0.2.4
  [d41bc354] NLSolversBase v7.8.3
  [2774e3e8] NLsolve v4.5.1
  [872c559c] NNlib v0.9.8
  [77ba4419] NaNMath v1.0.2
  [8913a72c] NonlinearSolve v2.8.2
  [d8793406] ObjectFile v0.4.1
  [6fe1bfb0] OffsetArrays v1.12.10
  [429524aa] Optim v1.7.8
  [3bd65402] Optimisers v0.3.1
  [7f7a1694] Optimization v3.19.3
  [36348300] OptimizationOptimJL v0.1.13
  [bac558e1] OrderedCollections v1.6.3
  [1dea7af3] OrdinaryDiffEq v6.59.2
  [90014a1f] PDMats v0.11.30
  [65ce6f38] PackageExtensionCompat v1.0.2
  [d96e819e] Parameters v0.12.3
  [06bb1623] PenaltyFunctions v0.3.0
  [e409e4f3] PoissonRandom v0.4.4
  [f517fe37] Polyester v0.7.9
  [1d0040c9] PolyesterWeave v0.2.1
  [85a6dd25] PositiveFactorizations v0.2.4
  [d236fae5] PreallocationTools v0.4.12
  [aea7be01] PrecompileTools v1.2.0
  [21216c6a] Preferences v1.4.1
  [27ebfcd6] Primes v0.5.5
  [33c8b6b6] ProgressLogging v0.1.4
  [92933f4c] ProgressMeter v1.9.0
  [1fd47b50] QuadGK v2.9.1
  [8a4e6c94] QuasiMonteCarlo v0.3.3
  [74087812] Random123 v1.6.1
  [e6cf234a] RandomNumbers v1.5.3
  [c1ae055f] RealDot v0.1.0
  [3cdcf5f2] RecipesBase v1.3.4
  [731186ca] RecursiveArrayTools v2.38.10
  [f2c3362d] RecursiveFactorization v0.2.21
  [189a3867] Reexport v1.2.2
  [ae029012] Requires v1.3.0
  [ae5879a3] ResettableStacks v1.1.1
  [37e2e3b7] ReverseDiff v1.15.1
  [79098fc4] Rmath v0.7.1
  [7e49a35a] RuntimeGeneratedFunctions v0.5.12
  [94e857df] SIMDTypes v0.1.0
  [476501e8] SLEEFPirates v0.6.42
  [0bca4576] SciMLBase v2.9.0
  [e9a6253c] SciMLNLSolve v0.1.9
  [c0aeaf25] SciMLOperators v0.3.7
  [1ed8b502] SciMLSensitivity v7.47.0
  [6c6a2e73] Scratch v1.2.1
  [efcf1570] Setfield v1.1.1
  [727e6d20] SimpleNonlinearSolve v0.1.25
  [699a6c99] SimpleTraits v0.9.4
  [ce78b400] SimpleUnPack v1.1.0
  [ed01d8cd] Sobol v1.5.0
  [a2af1166] SortingAlgorithms v1.2.0
  [47a9eef4] SparseDiffTools v2.13.0
  [dc90abb0] SparseInverseSubset v0.1.1
  [e56a9233] Sparspak v0.3.9
  [276daf66] SpecialFunctions v2.3.1
  [aedffcd0] Static v0.8.8
  [0d7ed370] StaticArrayInterface v1.4.1
  [90137ffa] StaticArrays v1.7.0
  [1e83bf80] StaticArraysCore v1.4.2
  [82ae8749] StatsAPI v1.7.0
  [2913bbd2] StatsBase v0.34.2
  [4c63d2b9] StatsFuns v1.3.0
⌅ [9672c7b4] SteadyStateDiffEq v1.16.1
  [789caeaf] StochasticDiffEq v6.63.2
  [7792a7ef] StrideArraysCore v0.5.1
  [09ab397b] StructArrays v0.6.16
  [53d494c1] StructIO v0.3.0
  [c3572dad] Sundials v4.20.1
  [2efcf032] SymbolicIndexingInterface v0.2.2
  [3783bdb8] TableTraits v1.0.1
  [bd369af6] Tables v1.11.1
  [5d786b92] TerminalLoggers v0.1.7
  [8290d209] ThreadingUtilities v0.5.2
  [a759f4b9] TimerOutputs v0.5.23
  [9f7883ad] Tracker v0.2.30
  [a2a6695c] TreeViews v0.3.0
  [d5829a12] TriangularSolve v0.1.20
  [410a4b4d] Tricks v0.1.8
  [781d530d] TruncatedStacktraces v1.4.0
  [3a884ed6] UnPack v1.0.2
  [013be700] UnsafeAtomics v0.2.1
  [d80eeb9a] UnsafeAtomicsLLVM v0.1.3
  [3d5dd08c] VectorizationBase v0.21.64
  [19fa3120] VertexSafeGraphs v0.2.0
  [e88e6eb3] Zygote v0.6.67
  [700de1a5] ZygoteRules v0.2.4
  [cd4c43a9] Dierckx_jll v0.1.0+0
⌅ [7cc45869] Enzyme_jll v0.0.89+0
  [1d5cc7b8] IntelOpenMP_jll v2023.2.0+0
  [dad2f222] LLVMExtra_jll v0.0.27+1
  [856f044c] MKL_jll v2023.2.0+0
  [efe28fd5] OpenSpecFun_jll v0.5.5+0
  [f50d1b31] Rmath_jll v0.4.0+0
⌅ [fb77eaff] Sundials_jll v5.2.1+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.3
  [76f85450] LibGit2
  [8f399da3] Libdl
  [37e2e46d] LinearAlgebra
  [56ddb016] Logging
  [d6f4376e] Markdown
  [a63ad114] Mmap
  [ca575930] NetworkOptions v1.2.0
  [44cfe95a] Pkg v1.9.2
  [de0858da] Printf
  [3fa0cd96] REPL
  [9a3f8284] Random
  [ea8e919c] SHA v0.7.0
  [9e88b42a] Serialization
  [1a1011a3] SharedArrays
  [6462fe0b] Sockets
  [2f01184e] SparseArrays
  [10745b16] Statistics v1.9.0
  [4607b0f0] SuiteSparse
  [fa267f1f] TOML v1.0.3
  [a4e569a6] Tar v1.10.0
  [8dfed614] Test
  [cf7118a7] UUIDs
  [4ec0a83e] Unicode
  [e66e0078] CompilerSupportLibraries_jll v1.0.5+0
  [deac9b47] LibCURL_jll v7.84.0+0
  [29816b5a] LibSSH2_jll v1.10.2+0
  [c8ffd9c3] MbedTLS_jll v2.28.2+0
  [14a3606d] MozillaCACerts_jll v2022.10.11
  [4536629a] OpenBLAS_jll v0.3.21+4
  [05823500] OpenLibm_jll v0.8.1+0
  [bea87d4a] SuiteSparse_jll v5.10.1+6
  [83775a58] Zlib_jll v1.2.13+0
  [8e850b90] libblastrampoline_jll v5.8.0+0
  [8e850ede] nghttp2_jll v1.48.0+0
  [3f19e933] p7zip_jll v17.4.0+0
Info Packages marked with ⌅ have new versions available but compatibility constraints restrict them from upgrading. To see why use `status --outdated -m`

Output of versioninfo()

Julia Version 1.9.2
Commit e4ee485e90 (2023-07-05 09:39 UTC)
Platform Info:
  OS: Linux (x86_64-redhat-linux)
  CPU: 24 × 12th Gen Intel(R) Core(TM) i9-12950HX
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-14.0.6 (ORCJIT, alderlake)
  Threads: 18 on 24 virtual cores
Environment:
  JULIA_NUM_THREADS = 18

Additional context

A little bit unrelated, but the model I'm trying to identify has 5 states and 38 parameters + 5*365=1825 parameters for starting points. Is it faster to use AutoZygote or AutoForwardDiff in this case? I'd think AutoForwardDiff is fast enough due to the small number of states but it seems slow and I have not been able to compare it with AutoZygote due to this bug.

SciML / DiffEqParamEstim.jl

AutoZygote incompatible with multiple shooting objective #229