How to use ChainRules to tell Zygote about a more specialized rule?

General Usage
, ,
1

I am trying to optimize an ODE problem specified with ModelingToolkit. The objective function indexes into the ODESolution by a symbol. There is a Zygote adjoint for getindex(::VectorOfArray, i), and VectorOfArray is an ancestor of ODESolution. It only handles integer indexing and I get errors when this pushback from ReverseArrayTools is called in the parameter optimization with a Num. (Thats when I started learning about AD, so I am not familiar with all the background)

I read about the ChainRules package and think that I should define a more specialized reverse rule for ODESolution or AbstractTimeseriesSolution that can handle this case.

Now, in my trials (excerpt below), this rule is never called (Julia 1.6 on Ubuntu), although
I patiently restarted the REPL in each try.

How are the rules selected? What do I need to read and do?

...  # setup of the problem, gradient is never called here
@variables t,x(t)
soltrue = solve(prob, Tsit5(), saveat = tsteps);

using ChainRulesCore
function ChainRulesCore.rrule(::typeof(getindex), VA::ODESolution, sym) 
  function ODESolution_getindex_pullback(Δ)
    @show Δ
    @show length(VA)
    @show VA
    @show VA.u
    # convert symbol to index
    i = issymbollike(sym) ? sym_to_index(sym, VA) : sym
    @show i
    # similar to VectorOfArray: return zero for non-matching indices
    Δ′ = [ (i == j ? Δ : zero(x)) for (x,j) in zip(VA.u, 1:length(VA))]
    (NO_FIELDS, Δ′)
  end  
  VA[sym], ODESolution_getindex_pullback(Δ)
end

  popt0 = [1.1]
  f1(p) = soltrue[x][1] * p[1] # simple case, reality use p to generate a modified solution
  f1(popt0)
  #using Zygote
  gr = Zygote.gradient(f1, popt0)

I still get the error of the pushback function of the rule defined in RecursiveArrayTools:

ZygoteRules.@adjoint function getindex(VA::AbstractVectorOfArray, i)
  function AbstractVectorOfArray_getindex_adjoint(Δ)
    @info "AbstractVectorOfArray_getindex_adjoint" # debug info inserted by me
    @show typeof(VA)
    Δ′ = [ (i == j ? Δ : zero(x)) for (x,j) in zip(VA.u, 1:length(VA))]
    (Δ′,nothing)
  end
  VA[i],AbstractVectorOfArray_getindex_adjoint
end


[ Info: AbstractVectorOfArray_getindex_adjoint
typeof(VA) = ODESolution{Float64, 2, ...
...
ERROR: TypeError: non-boolean (Num) used in boolean context
Stacktrace:
  [1] (::RecursiveArrayTools.var"#80#82"{Vector{Float64}, Num})(::Tuple{Vector{Float64}, Int64})
    @ RecursiveArrayTools ./none:0
...
  [6] Pullback
    @ ./REPL[21]:1 [inlined]
  [7] (::typeof(∂(f1)))(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/RxTZu/src/compiler/interface2.jl:0
  [8] (::Zygote.var"#41#42"{typeof(∂(f1))})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/RxTZu/src/compiler/interface.jl:41
  [9] gradient(f::Function, args::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/RxTZu/src/compiler/interface.jl:59
 [10] top-level scope
    @ REPL[24]:1
2

Can you share your test case? This is hard to debug without it. Open an issue on DifferentialEquations.jl about this with the test case. I am not sure this will be an easy one.

3

I opened an Issue on DifferentialEquations.jl. (Accidentally this went to OrdinaryDiffEq.jl, can one transfer the issue to the correct repo?)