Performance issue: differentiating towards a subset of parameters

[ArnoStrouwen]

using DiffEqSensitivity
using OrdinaryDiffEq

using Zygote
using ForwardDiff

using BenchmarkTools

function fiip(du,u,p,t)
    du[1] = dx = p[1]*u[1] - p[2]*u[1]*u[2]
    du[2] = dy = -p[3]*u[2] + p[4]*u[1]*u[2]
end
u0 = [1.0,1.0]
p = [1.5,1.0,3.0,1.0]
prob = ODEProblem(fiip,u0,(0.0,10.0),[1.5,1.0,3.0,1.0],reltol = 1e-3) 

function cost(p)
    _prob = remake(prob,p=p)
    sol = solve(_prob,Tsit5(),sensealg=ForwardSensitivity(),saveat=0.1)
    sum(sol)
end
@btime ForwardDiff.gradient(cost,p) # 69.159 μs (418 allocations: 43.48 KiB)
@btime Zygote.gradient(cost,p) # 320.949 μs (1323 allocations: 137.23 KiB)

function cost(p)
    _prob = remake(prob,p=vcat(p,zeros(100))) # add useless zeros
    sol = solve(_prob,Tsit5(),sensealg=ForwardSensitivity(),saveat=0.1)
    sum(sol)
end
# barely slower:
@btime ForwardDiff.gradient(cost,p) # 74.737 μs (420 allocations: 48.55 KiB)
# severe slowdown
@btime Zygote.gradient(cost,p) # 2.576 ms (1336 allocations: 732.30 KiB)

[ChrisRackauckas]

This is because in reverse mode the Jacobian has to be constructed from the forward mode which grows. I don't think that's going to be the right strategy and instead we need to tie into the real forward mode of an AD with both to avoid that.