DimensionMismatch error

[dpo]

I asked this question in the #autodiff channel on Slack but haven't found out if I'm doing something wrong, hitting a bug in ReverseDiff, or hitting an incompatibility with DifferentialEquations.

I’m trying to differentiate through a system of ODEs. The function F returns simulations F(x) and I’m trying to obtain the product J(x)*u and J(x)ᵀ*v where J is the Jacobian of F. It works with ForwardDiff, but ReverseDiff throws an error…

using DifferentialEquations, DiffEqSensitivity, ForwardDiff, ReverseDiff

function ODE(dx, x, p, t)
  V, W = x
  I, μ, a, b, c = p
  dx[1] = (V - V^3 / 3 - W + I) / μ
  dx[2] = μ * (a * V - b * W + c)
end

prob = ODEProblem(ODE, [2.0; 0.0], (0.0, 20.0), [0.5, 0.08, 1.0, 0.8, 0.7])

function F(x)
  temp_prob = remake(prob, p = x)
  sol = solve(temp_prob, Vern9(), saveat = 0.2)
  return vec(sol)
end

jprod_fwd(f, x, u) = ForwardDiff.derivative(t -> f(x + t * u), 0)
jprod_fwd(F, rand(5), rand(5))  # works

jprod_rev(f, x, u) = ReverseDiff.jacobian(t -> f(x + t[1] * u), [zero(eltype(x))])
jprod_rev(F, rand(5), rand(5))  # throws DimensionMismatch ?!

jtprod_rev(f, x, u) = ReverseDiff.gradient(z -> dot(f(z), u), x)
jtprod_rev (generic function with 1 method)

jtprod_rev(F, rand(5), rand(202))
ERROR: DimensionMismatch("arrays could not be broadcast to a common size; got a dimension with lengths 5 and 6")

Here is the full error message: https://gist.github.com/7544d5c680995d4634cab19a2cdccfd6

pkg> status
      Status `~/dev/julia/BLA/myenv/Project.toml`
  [41bf760c] DiffEqSensitivity v6.57.0
  [0c46a032] DifferentialEquations v6.18.0
  [f6369f11] ForwardDiff v0.10.19
  [37e2e3b7] ReverseDiff v1.9.0

Many thanks in advance!

ps: I am aware that using ReverseDiff for J(x)*u isn't the most efficient in many cases, but it should work. It works with Zygote.