# Help on ForwardDiff and Zygote for ODEProblem with mass matrix

**URL:** <https://discourse.julialang.org/t/help-on-forwarddiff-and-zygote-for-odeproblem-with-mass-matrix/117857>\
**Category:** Modelling & Simulations\
**Tags:** question, sciml, autodiff, differentialequation\
**Created:** [August 5, 2024, 7:57pm UTC](https://discourse.julialang.org/t/help-on-forwarddiff-and-zygote-for-odeproblem-with-mass-matrix/117857 "2024-08-05T19:57:18Z")\
**Posts on this page:** 9\
**Page:** 1

<div class="post-metadata">

**Author:** ![ccz](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ccz/32/211119_2.png) [@ccz](https://discourse.julialang.org/u/ccz)\
**Post date:** [August 5, 2024, 7:57pm UTC](https://discourse.julialang.org/t/help-on-forwarddiff-and-zygote-for-odeproblem-with-mass-matrix/117857/1 "2024-08-05T19:57:19Z")

</div>

Hi there,

I’m new to `SciMLSensitivity.jl` and `Zygote.jl` and have been playing with these two libraries for a toy ODE system with a mass matrix. However, the Jacobians produced by the two libraries are different, and I haven’t been able to figure out why.

Could you please take a look at my minimal code example and give me a hint on what might be causing this issue?

Regards,  
Charlie

```julia
using DifferentialEquations
using SciMLSensitivity
using Zygote
using ForwardDiff

K = [1.0 0 0
     0 1.0 0
     0 0 1.0]
function dudt(du, u, p, t)
    du .= K*u
    # nothing
end
M = [1.0 0 0
     0 1.0 0
     0 0 1.0]
f = ODEFunction(dudt, mass_matrix = M)

# target (vector) function
function forward_vector_target(u)
    prob = ODEProblem(f, u, (0.0, 1.0), [1.0, 1.0, 1.0])
    sol = solve(prob, Rodas5(), reltol = 1e-8, abstol = 1e-8)
    return sol[end]
end

u0 = [1.0, 1.0, 1.0]

# 1. reserse mode
rs_jac = Zygote.jacobian(forward_vector_target, u0)[1]
println(rs_jac)
# output:
# 3×3 Matrix{Float64}:
# 2.71828 0.0 0.0
# 2.71828 2.71828 0.0
# 2.71828 2.71828 2.71828

# 2. forward mode
fs_jac = ForwardDiff.jacobian(forward_vector_target, u0)
println(fs_jac)
# output:
# 3×3 Matrix{Float64}:
# 2.71828 0.0 0.0
# 0.0 2.71828 0.0
# 0.0 0.0 2.71828

```

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [August 5, 2024, 10:08pm UTC](https://discourse.julialang.org/t/help-on-forwarddiff-and-zygote-for-odeproblem-with-mass-matrix/117857/2 "2024-08-05T22:08:51Z")

</div>

Here’s a very rough experiment which points towards the ForwardDiff one being correct:

```julia
julia> y0 = forward_vector_target(u0)
3-element Vector{Float64}:
 2.7182818287174166
 2.7182818287174166
 2.7182818287174166

julia> forward_vector_target(u0 .+ [1,0,0] ./ 10) .- y0
3-element Vector{Float64}:
  0.2718281828685605
 -2.8883562208648073e-12
 -2.8883562208648073e-12

julia> forward_vector_target(u0 .+ [0,1,0] ./ 10) .- y0
3-element Vector{Float64}:
 -2.8883562208648073e-12
  0.2718281828685605
 -2.8883562208648073e-12

julia> forward_vector_target(u0 .+ [0,0,1] ./ 10) .- y0
3-element Vector{Float64}:
 -2.8883562208648073e-12
 -2.8883562208648073e-12
  0.2718281828685605

julia> fs_jac = ForwardDiff.jacobian(forward_vector_target, u0)
3×3 Matrix{Float64}:
 2.71828 0.0 0.0
 0.0 2.71828 0.0
 0.0 0.0 2.71828

```

(I wondered whether a known bug in ForwardDiff v0.10 might be the problem, so I tried with ForwardDiff v0.11-DEV to check… but both versions give the same.)

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [August 5, 2024, 11:49pm UTC](https://discourse.julialang.org/t/help-on-forwarddiff-and-zygote-for-odeproblem-with-mass-matrix/117857/3 "2024-08-05T23:49:05Z")

</div>

It looks like it’s a Zygote thing. When I try to recreate the Jacobian it’s fine:

```julia
julia> [Zygote.pullback(forward_vector_target, u0)[2]([1.0,0.0,0.0])[1]';
        Zygote.pullback(forward_vector_target, u0)[2]([0.0,1.0,0.0])[1]';
        Zygote.pullback(forward_vector_target, u0)[2]([0.0,0.0,1.0])[1]']
3×3 Matrix{Float64}:
 2.71828 0.0 0.0
 0.0 2.71828 0.0
 0.0 0.0 2.71828

```

That’s pulling back the basis vectors, and it seems fine. That’s what I would’ve assumed Zygote.jacobian was doing, so I’m surprised it’s a different answer. @mcabbott do you know what Zygote.jacobian actually calculating with?

```julia
julia> [Zygote.gradient((u)->forward_vector_target(u)[1], u0)[1]';
       Zygote.gradient((u)->forward_vector_target(u)[2], u0)[1]';
       Zygote.gradient((u)->forward_vector_target(u)[3], u0)[1]']
3×3 Matrix{Float64}:
 2.71828 0.0 0.0
 0.0 2.71828 0.0
 0.0 0.0 2.71828

```

Gradients are also fine, so this is very specific to the Zygote.jacobian interface, and I don’t know its implementation.

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [August 5, 2024, 11:55pm UTC](https://discourse.julialang.org/t/help-on-forwarddiff-and-zygote-for-odeproblem-with-mass-matrix/117857/4 "2024-08-05T23:55:36Z")

</div>

I didn’t manage to run the Zygote version locally.

But the difference between what you write and `Zygote.jacobian` is that you run the forward pass 3 times, while it is something more like this (untested):

```julia
y, bk = Zygote.pullback(forward_vector_target, u0)
hcat(bk([1,0,0.]), bk([0,1,0.]), bk([0,0,1.]))

```

i.e. saving and re-using the pullback function.

If this snippet differs from what you write, then I believe the rule you are providing Zygote is invalid, e.g. it is overwriting some buffer captured by `bk` such that it does not work on the second call.

(Overwriting such buffers does allow more efficient rules, if the pullback is certain to only be used once. It’s easy to forget to test for this. At some point, there was work towards adding a flag for this to ChainRulesCore, and making Zygote use it, but IIRC this never got merged. Edit, see [this CRC PR](https://github.com/JuliaDiff/ChainRulesCore.jl/pull/592).)

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [August 6, 2024, 1:27am UTC](https://discourse.julialang.org/t/help-on-forwarddiff-and-zygote-for-odeproblem-with-mass-matrix/117857/5 "2024-08-06T01:27:17Z")

</div>

Okay, that narrows it down. Here’s a leaner MWE, it actually doesn’t involve mass matrices at all:

```julia
using OrdinaryDiffEq
using SciMLSensitivity
using Zygote
using ForwardDiff

function dudt(du, u, p, t)
    du .= u
    nothing
end

# target (vector) function
function forward_vector_target(u)
    prob = ODEProblem(dudt, u, (0.0, 1.0), [1.0, 1.0, 1.0])
    sol = solve(prob, FBDF(), reltol = 1e-8, abstol = 1e-8, sensealg = ForwardDiffSensitivity())
    return sol[end]
end

y, back = Zygote.pullback(forward_vector_target, u0);
[back([1.0,0.0,0.0])[1]';
back([0.0,1.0,0.0])[1]';
back([0.0,0.0,1.0])[1]']

```

```julia
julia> y, back = Zygote.pullback(forward_vector_target, u0);

julia> [back([1.0,0.0,0.0])[1]';
       back([0.0,1.0,0.0])[1]';
       back([0.0,0.0,1.0])[1]']
unthunk(du0) = [2.7182818436166185, 0.0, 0.0]
unthunk(du0) = [0.0, 2.7182818436166185, 0.0]
unthunk(du0) = [0.0, 0.0, 2.7182818436166185]
3×3 Matrix{Float64}:
 2.71828 0.0 0.0
 2.71828 2.71828 0.0
 2.71828 2.71828 2.71828

```

I added a printout of the `du0` computed by the adjoint method. That confirms that SciMLSensitivity is computing the correct value each time, so the adjoint system is completely fine with being reused with new values. So this means I’m handing

```julia
            (NoTangent(), NoTangent(), NoTangent(), du0, dp, NoTangent(),
                ntuple(_ -> NoTangent(), length(args))...)

```

back to Zygote and that `du0` is correct. The next thing above it is the `ODEProblem`, so I check:

```julia
function ChainRulesCore.rrule(::Type{ODEProblem}, args...; kwargs...)
    function ODEProblemAdjoint(ȳ)
        @show ȳ.u0
        (NoTangent(), ȳ.f, ȳ.u0, ȳ.tspan, ȳ.p, ȳ.kwargs, ȳ.problem_type)
    end

    ODEProblem(args...; kwargs...), ODEProblemAdjoint
end

```

```julia
unthunk(du0) = [2.7182818436166185, 0.0, 0.0]
ȳ.u0 = [2.7182818436166185, 0.0, 0.0]
unthunk(du0) = [0.0, 2.7182818436166185, 0.0]
ȳ.u0 = [2.7182818436166185, 2.7182818436166185, 0.0]
unthunk(du0) = [0.0, 0.0, 2.7182818436166185]
ȳ.u0 = [2.7182818436166185, 2.7182818436166185, 2.7182818436166185]
3×3 Matrix{Float64}:
 2.71828 0.0 0.0
 2.71828 2.71828 0.0
 2.71828 2.71828 2.71828

```

so the buffer that Zygote gives it has already merged the runs.

So since the pullback of `solve` is correct, what is changing the value before it ends up in the ODEProblem adjoint? That surely looks like Zygote bug then, since it should just be handling that same `du0` back up to the ODEProblem adjoint, which is then just the identity function derivative?

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [August 6, 2024, 3:22am UTC](https://discourse.julialang.org/t/help-on-forwarddiff-and-zygote-for-odeproblem-with-mass-matrix/117857/6 "2024-08-06T03:22:04Z")

</div>

Something closed over by `back` is certainly being mutated. For example, here the value of `back([0,1,0.])` is inconsistent:

```julia
julia> y, back = Zygote.pullback(forward_vector_target, u0);

julia> du1, = back([1,0,0.])
([2.7182818436166185, 0.0, 0.0],)

julia> back([0,1,0.]) # second call
([2.7182818436166185, 2.7182818436166185, 0.0],)

julia> back([0,1,0.]) # third call, with same input as second
([2.7182818436166185, 5.436563687233237, 0.0],)

julia> (du1,) # check the return wasn't mutated after printing
([2.7182818436166185, 0.0, 0.0],)

julia> y, back = Zygote.pullback(forward_vector_target, u0);

julia> back([0,1,0.]) # first call
([0.0, 2.7182818436166185, 0.0],)

```

The returned arrays here have distinct `objectid`s, so it isn’t literally mutating the returned array.

What exactly is printing `unthunk(du0)` you haven’t shown. But I agree that stage looks correct. Surely the straight-through `rrule` you show isn’t introducing any mutable state.

Note that ChainRulesTestUtils.jl tests for this problem, so assuming you define `rrule`s, you can try to make a Zygote-free reproducer:

```julia
julia> using Zygote, ChainRulesCore, ChainRulesTestUtils

julia> sq1(x) = x .* x;

julia> sq2(x) = x .* x; # same function, but will give it a bad rrule

julia> function ChainRulesCore.rrule(::typeof(sq1), x)
         sq1(x), dy -> (NoTangent(), 2 .* x .* dy) # good (at least for real numbers)
       end

julia> function ChainRulesCore.rrule(::typeof(sq2), x)
         BUF = zero.(x)
         return sq2(x), dy -> (NoTangent(), copy(BUF .+= 2 .* x .* dy)) # bad!
       end

julia> Zygote.gradient(x -> sum(x.*x), [1,2,3.])
([2.0, 4.0, 6.0],)

julia> Zygote.gradient(sum∘sq1, [1,2,3.])
([2.0, 4.0, 6.0],)

julia> Zygote.gradient(sum∘sq2, [1,2,3.]) # ok
([2.0, 4.0, 6.0],)

julia> Zygote.jacobian(x -> x.*x, [1,2,3.])
([2.0 0.0 0.0; 0.0 4.0 0.0; 0.0 0.0 6.0],)

julia> Zygote.jacobian(sq1, [1,2,3.])
([2.0 0.0 0.0; 0.0 4.0 0.0; 0.0 0.0 6.0],)

julia> Zygote.jacobian(sq2, [1,2,3.]) # wrong!
([2.0 0.0 0.0; 2.0 4.0 0.0; 2.0 4.0 6.0],)

julia> _, back = Zygote.pullback(sq2, [1,2,3.]);

julia> back([1,0,0.])
([2.0, 0.0, 0.0],)

julia> back([1,0,0.]) # second call, different answer
([4.0, 0.0, 0.0],)

julia> test_rrule(sq1, [1,2,3.])
Test Summary: | Pass Total Time
test_rrule: sq1 on Vector{Float64} | 7 7 0.0s
Test.DefaultTestSet("test_rrule: sq1 on Vector{Float64}", Any[], 7, false, false, true, 1.72291392558364e9, 1.7229139256149e9, false, "/Users/me/.julia/packages/ChainRulesTestUtils/Ko1Wr/src/testers.jl")

julia> test_rrule(sq2, [1,2,3.])
test_rrule: sq2 on Vector{Float64}: Test Failed at /Users/me/.julia/packages/ChainRulesTestUtils/Ko1Wr/src/check_result.jl:24
  Expression: isapprox(actual, expected; kwargs...)
  Problem: cotangent for input 2, 3-element Vector{Float64}
   Evaluated: isapprox([14.12, -46.72, 77.88], [7.059999999999687, -23.36000000000044, 38.94000000000155]; rtol = 1.0e-9, atol = 1.0e-9)

```

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [August 6, 2024, 7:37am UTC](https://discourse.julialang.org/t/help-on-forwarddiff-and-zygote-for-odeproblem-with-mass-matrix/117857/7 "2024-08-06T07:37:58Z")

</div>

When I isolate the SciMLSensitivity part, I can confirm it’s fine:

```julia
prob = ODEProblem(dudt, u0, (0.0, 1.0), [1.0, 1.0, 1.0])
ff(x) = DiffEqBase.solve_up(prob, ForwardDiffSensitivity(), x, prob.p, FBDF())[end]

y, back = Zygote.pullback(ff, u0);
[back([1.0,0.0,0.0])[1]';
back([0.0,1.0,0.0])[1]';
back([0.0,0.0,1.0])[1]']

```

```julia
unthunk(du0) = [2.7196609758646866, 0.0, 0.0]
unthunk(du0) = [0.0, 2.7196609758646866, 0.0]
unthunk(du0) = [0.0, 0.0, 2.7196609758646866]
3×3 Matrix{Float64}:
 2.71966 0.0 0.0
 0.0 2.71966 0.0
 0.0 0.0 2.71966

```

Isolating it to the ODEProblem overload:

```julia
julia> g(x) = ODEProblem(dudt, x, (0.0, 1.0), [1.0, 1.0, 1.0]).u0
g (generic function with 1 method)

julia> y, back = Zygote.pullback(g, u0);

julia> [back([1.0,0.0,0.0])[1]';
       back([0.0,1.0,0.0])[1]';
       back([0.0,0.0,1.0])[1]']
ȳ.u0 = [1.0, 0.0, 0.0]
ȳ.u0 = [1.0, 1.0, 0.0]
ȳ.u0 = [1.0, 1.0, 1.0]
3×3 Matrix{Float64}:
 1.0 0.0 0.0
 1.0 1.0 0.0
 1.0 1.0 1.0

```

So it is this overload:

```julia
function ChainRulesCore.rrule(::Type{ODEProblem}, args...; kwargs...)
    function ODEProblemAdjoint(ȳ)
        (NoTangent(), ȳ.f, ȳ.u0, ȳ.tspan, ȳ.p, ȳ.kwargs, ȳ.problem_type)
    end
    ODEProblem(args...; kwargs...), ODEProblemAdjoint
end

```

since that’s the only one involved on our end. So then here’s the new MWE:

```julia
using SciMLBase, Zygote
u0 = [1.0, 1.0, 1.0]
g(x) = ODEProblem((du,u,p,t)->du.=u, x, (0.0, 1.0), [1.0, 1.0, 1.0]).u0
y, back = Zygote.pullback(g, u0);
back([1.0,0.0,0.0])[1]
back([1.0,0.0,0.0])[1]

```

```julia

julia> back([1.0,0.0,0.0])[1]
3-element Vector{Float64}:
 1.0
 0.0
 0.0

julia> back([1.0,0.0,0.0])[1]
3-element Vector{Float64}:
 2.0
 0.0
 0.0

```

Note that `ȳ.u0` is already wrong, so that means it’s not the SciML overload. There’s only one function left then, which is the `getproperty(prob, :u0)`.

The conclusion then is that the problem must be the Zygote literal\_getproperty fallback?

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [August 6, 2024, 7:43am UTC](https://discourse.julialang.org/t/help-on-forwarddiff-and-zygote-for-odeproblem-with-mass-matrix/117857/8 "2024-08-06T07:43:51Z")

</div>

Indeed, I isolated it to Zygote’s getproperty fallback by simply defining a literal getproperty overload. Once that’s done, it goes away:

```julia
using SciMLBase, Zygote
u0 = [1.0, 1.0, 1.0]; odef(du,u,p,t)=nothing;
g(x) = ODEProblem(odef, x, (0.0, 1.0), [1.0, 1.0, 1.0]).u0
y, back = Zygote.pullback(g, u0);
back([1.0,0.0,0.0])[1]
back([1.0,0.0,0.0])[1]

Zygote.@adjoint function Zygote.literal_getproperty(prob::ODEProblem,
        ::Val{:u0})
    prob.u0, p̄ -> (ODEProblem(prob.f, p̄, prob.tspan, prob.p),)
end

y, back = Zygote.pullback(g, u0);
back([1.0,0.0,0.0])[1]
back([1.0,0.0,0.0])[1]

```

```julia
julia> back([1.0,0.0,0.0])[1]
3-element Vector{Float64}:
 1.0
 0.0
 0.0

julia> back([1.0,0.0,0.0])[1]
3-element Vector{Float64}:
 2.0
 0.0
 0.0

julia> Zygote.@adjoint function Zygote.literal_getproperty(prob::ODEProblem,
               ::Val{:u0})
           prob.u0, p̄ -> (ODEProblem(prob.f, p̄, prob.tspan, prob.p),)
       end

julia> y, back = Zygote.pullback(g, u0);

julia> back([1.0,0.0,0.0])[1]
3-element Vector{Float64}:
 1.0
 0.0
 0.0

julia> back([1.0,0.0,0.0])[1]
3-element Vector{Float64}:
 1.0
 0.0
 0.0

```

@mcabbott is there a mutation in that getproperty fallback?

---

<div class="post-metadata">

**Author:** ![ccz](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ccz/32/211119_2.png) [@ccz](https://discourse.julialang.org/u/ccz)\
**Post date:** [August 6, 2024, 8:27am UTC](https://discourse.julialang.org/t/help-on-forwarddiff-and-zygote-for-odeproblem-with-mass-matrix/117857/9 "2024-08-06T08:27:29Z")

</div>

Thanks very much for digging so deep into this problem!

I want to add that if I 1) define the ODEProblem out of the target function and 2) enforce its parameter `u0` with `remake` inside the function, then the pullback gives the correct jacobian. I’m not sure if this tiny observation is relevant at all, but I’m pasting the code (based on your MWE) here:

```julia
using OrdinaryDiffEq
using SciMLSensitivity
using Zygote
using ForwardDiff

function dudt(du, u, p, t)
    du .= u
    nothing
end

u0 = [1.0, 1.0, 1.0]
prob = ODEProblem(dudt, u0, (0.0, 1.0), [1.0, 1.0, 1.0])

# target (vector) function
function forward_vector_target(u)
    _prob = remake(prob, u0 = u)
    sol = solve(_prob, FBDF(), reltol = 1e-8, abstol = 1e-8, sensealg = ForwardDiffSensitivity())
    return sol[end]
end

y, back = Zygote.pullback(forward_vector_target, u0);
[back([1.0,0.0,0.0])[1]';
back([0.0,1.0,0.0])[1]';
back([0.0,0.0,1.0])[1]']

```