Finding mutations that Zygote complains about

Specific Domains
, , ,
1

Hey,

Apparently, I have some hidden mutations in my code that I cannot find. Unfortunately, Zygote easily finds them and complains about it.
I get call stacks 40+ layers deep, the last of my own functions being near the top (layer 6):

  [1] error(s::String)
    @ Base ./error.jl:44
  [2] _throw_mutation_error(f::Function, args::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/lib/array.jl:70
  [3] (::Zygote.var"#713#714"{Vector{Float64}})(::Nothing)
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/lib/array.jl:85
  [4] (::Zygote.var"#715#716"{Zygote.var"#713#714"{Vector{Float64}}})(Δ::Nothing)
    @ Zygote ~/.julia/packages/ZygoteRules/CkVIK/src/adjoint.jl:72
  [5] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:100
  [6] rhs_jac!
    @ ~/Projects/NODE/src/Model.jl:441 [inlined]
  [7] (::Zygote.Pullback{Tuple{…}, Any})(Δ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:0
  [8] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:100
  [9] (::Zygote.Pullback{Tuple{…}, Any})(Δ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:100
 [10] #_vecjacobian!##16
    @ ~/.julia/packages/SciMLSensitivity/3s0Zy/src/derivative_wrappers.jl:659 [inlined]
 [11] (::Zygote.Pullback{Tuple{…}, Any})(Δ::SubArray{Float64, 1, Vector{…}, Tuple{…}, true})
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:0
 [12] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{…}})(Δ::SubArray{Float64, 1, Vector{…}, Tuple{…}, true})
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface.jl:97
 [13] _vecjacobian!(dλ::SubArray{…}, y::Vector{…}, λ::SubArray{…}, p::ComponentArrays.ComponentVector{…}, t::Float64, S::SciMLSensitivity.ODEInterpolatingAdjointSensitivityFunction{…}, isautojacvec::SciMLSensitivity.ZygoteVJP, dgrad::SubArray{…}, dy::Nothing, W::Nothing)
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/3s0Zy/src/derivative_wrappers.jl:630
 [14] #vecjacobian!#19
    @ ~/.julia/packages/SciMLSensitivity/3s0Zy/src/derivative_wrappers.jl:257 [inlined]
 [15] vecjacobian!
    @ ~/.julia/packages/SciMLSensitivity/3s0Zy/src/derivative_wrappers.jl:252 [inlined]
 [16] (::SciMLSensitivity.ODEInterpolatingAdjointSensitivityFunction{…})(du::Vector{…}, u::Vector{…}, p::ComponentArrays.ComponentVector{…}, t::Float64)
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/3s0Zy/src/interpolating_adjoint.jl:153
...

Background: This is part of a Neural ODE project, where I try to use adjoints for gradient computation.

How can I systematically debug / work through the code to find those mutations?

Thanks,
Neodym

2

It’s this function right here. You can pull it out and directly do the vjp call in order to recreate it outside of the SciMLSensitivity context, and that’ll make it easier to debug.

But, I’d recommend just trying MooncakeVJP or EnzymeVJP these days. That should just be faster too. Are you directly setting ZygoteVJP?