Unexpected error when DataInterpolations.jl updates from 5.0.0 to 5.3.0.

[chooron]

Describe the bug 🐞

Hello, I am using Zygote.jl to solve gradients with DataInterpolations.jl, but an error occurred.

Expected behavior

Expected to run successfully.

Minimal Reproducible Example 👇

#* test for datainterpolation and zygote
using DataInterpolations
using SciMLSensitivity, OrdinaryDiffEq, Zygote

function loss(u0, p)
    itp = LinearInterpolation(rand(length(0.0:0.1:10.0)), 0.0:0.1:10.0)
    function fiip(du, u, p, t)
        du[1] = dx = p[1] * u[1] - p[2] * u[1] * u[2] + itp(t)
        du[2] = dy = -p[3] * u[2] + p[4] * u[1] * u[2] + itp(t)
    end

    prob = ODEProblem(fiip, u0, (0.0, 10.0), p)
    sum(solve(prob, Tsit5(), u0=u0, p=p, saveat=0.1))
end

p = [1.5, 1.0, 3.0, 1.0]
u0 = [1.0; 1.0]
du0, dp = Zygote.gradient(loss, u0, p)

Error & Stacktrace ⚠️

ERROR: ArgumentError: invalid index: Dual{Nothing}(2,0,0,1) of type ForwardDiff.Dual{Nothing, Int64, 3}
Stacktrace:
  [1] to_index(i::ForwardDiff.Dual{Nothing, Int64, 3})
    @ Base .\indices.jl:300
  [2] to_index(A::Vector{Float64}, i::ForwardDiff.Dual{Nothing, Int64, 3})
    @ Base .\indices.jl:277
  [3] _to_indices1(A::Vector{Float64}, inds::Tuple{Base.OneTo{Int64}}, I1::ForwardDiff.Dual{Nothing, Int64, 3})
    @ Base .\indices.jl:359
  [4] to_indices
    @ .\indices.jl:354 [inlined]
  [5] to_indices
    @ .\indices.jl:345 [inlined]
  [6] getindex
    @ .\abstractarray.jl:1288 [inlined]
  [7] getindex
    @ D:\Julia\Julia-1.10.0\packages\packages\ReadOnlyArrays\I7g0C\src\ReadOnlyArrays.jl:63 [inlined]
  [8] linear_interpolation_parameters(u::ReadOnlyArrays.ReadOnlyVector{…}, t::ReadOnlyArrays.ReadOnlyVector{…}, idx::ForwardDiff.Dual{…})
    @ DataInterpolations D:\Julia\Julia-1.10.0\packages\packages\DataInterpolations\Bg3yw\src\parameter_caches.jl:11
  [9] (::Zygote.var"#1388#1389"{…})(::ReadOnlyArrays.ReadOnlyVector{…}, ::ReadOnlyArrays.ReadOnlyVector{…}, ::Int64)
    @ Zygote D:\Julia\Julia-1.10.0\packages\packages\Zygote\nsBv0\src\lib\broadcast.jl:276
 [10] _broadcast_getindex_evalf
    @ Base.Broadcast .\broadcast.jl:709 [inlined]
 [11] _broadcast_getindex
    @ Base.Broadcast .\broadcast.jl:682 [inlined]
 [12] getindex
    @ Base.Broadcast .\broadcast.jl:636 [inlined]
 [13] copy
    @ Base.Broadcast .\broadcast.jl:942 [inlined]
 [14] materialize
    @ Base.Broadcast .\broadcast.jl:903 [inlined]
 [15] broadcast_forward
    @ Zygote D:\Julia\Julia-1.10.0\packages\packages\Zygote\nsBv0\src\lib\broadcast.jl:282 [inlined]
 [16] _broadcast_generic
    @ Zygote D:\Julia\Julia-1.10.0\packages\packages\Zygote\nsBv0\src\lib\broadcast.jl:212 [inlined]
 [17] adjoint
    @ Zygote D:\Julia\Julia-1.10.0\packages\packages\Zygote\nsBv0\src\lib\broadcast.jl:205 [inlined]
 [18] _pullback(::Zygote.Context{…}, ::typeof(Base.Broadcast.broadcasted), ::Base.Broadcast.DefaultArrayStyle{…}, ::typeof(DataInterpolations.linear_interpolation_parameters), ::Base.RefValue{…}, ::Base.RefValue{…}, ::UnitRange{…})
    @ Zygote D:\Julia\Julia-1.10.0\packages\packages\ZygoteRules\M4xmc\src\adjoint.jl:67
 [19] _apply(::Function, ::Vararg{Any})
    @ Core .\boot.jl:838
 [20] adjoint
    @ D:\Julia\Julia-1.10.0\packages\packages\Zygote\nsBv0\src\lib\lib.jl:203 [inlined]
 [21] _pullback
    @ D:\Julia\Julia-1.10.0\packages\packages\ZygoteRules\M4xmc\src\adjoint.jl:67 [inlined]
 [22] broadcasted
    @ .\broadcast.jl:1347 [inlined]
 [23] LinearParameterCache
    @ D:\Julia\Julia-1.10.0\packages\packages\DataInterpolations\Bg3yw\src\parameter_caches.jl:6 [inlined]
 [24] _pullback(::Zygote.Context{…}, ::Type{…}, ::ReadOnlyArrays.ReadOnlyVector{…}, ::ReadOnlyArrays.ReadOnlyVector{…})
    @ Zygote D:\Julia\Julia-1.10.0\packages\packages\Zygote\nsBv0\src\compiler\interface2.jl:0
 [25] #LinearInterpolation#1
    @ D:\Julia\Julia-1.10.0\packages\packages\DataInterpolations\Bg3yw\src\interpolation_caches.jl:33 [inlined]
 [26] _pullback(::Zygote.Context{…}, ::DataInterpolations.var"##LinearInterpolation#1", ::Bool, ::Bool, ::Type{…}, ::Vector{…}, ::StepRangeLen{…})
    @ Zygote D:\Julia\Julia-1.10.0\packages\packages\Zygote\nsBv0\src\compiler\interface2.jl:0
 [27] LinearInterpolation
    @ D:\Julia\Julia-1.10.0\packages\packages\DataInterpolations\Bg3yw\src\interpolation_caches.jl:31 [inlined]
 [28] _pullback(::Zygote.Context{…}, ::Type{…}, ::Vector{…}, ::StepRangeLen{…})
    @ Zygote D:\Julia\Julia-1.10.0\packages\packages\Zygote\nsBv0\src\compiler\interface2.jl:0
 [29] loss
    @ e:\JlCode\JuliaEnvTest\src\datainterpolation_zygote.jl:6 [inlined]
 [30] _pullback(::Zygote.Context{false}, ::typeof(loss), ::Vector{Float64}, ::Vector{Float64})
    @ Zygote D:\Julia\Julia-1.10.0\packages\packages\Zygote\nsBv0\src\compiler\interface2.jl:0
 [31] pullback(::Function, ::Zygote.Context{false}, ::Vector{Float64}, ::Vararg{Vector{Float64}})
    @ Zygote D:\Julia\Julia-1.10.0\packages\packages\Zygote\nsBv0\src\compiler\interface.jl:90
 [32] pullback(::Function, ::Vector{Float64}, ::Vector{Float64})
    @ Zygote D:\Julia\Julia-1.10.0\packages\packages\Zygote\nsBv0\src\compiler\interface.jl:88
 [33] gradient(::Function, ::Vector{Float64}, ::Vararg{Vector{Float64}})
    @ Zygote D:\Julia\Julia-1.10.0\packages\packages\Zygote\nsBv0\src\compiler\interface.jl:147
 [34] top-level scope
    @ e:\JlCode\JuliaEnvTest\src\datainterpolation_zygote.jl:18
Some type information was truncated. Use `show(err)` to see complete types.

Environment (please complete the following information):

• Output of using Pkg; Pkg.status()

Status `E:\JlCode\JuliaEnvTest\Project.toml`
  [82cc6244] DataInterpolations v5.3.0
  [1ed8b502] SciMLSensitivity v7.64.0
  [e88e6eb3] Zygote v0.6.70

• Output of versioninfo()

Julia Version 1.10.0
Commit 3120989f39 (2023-12-25 18:01 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Windows (x86_64-w64-mingw32)
  CPU: 24 × 12th Gen Intel(R) Core(TM) i9-12900HX
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-15.0.7 (ORCJIT, alderlake)
  Threads: 1 on 24 virtual cores
Environment:
  JULIA_DEPOT_PATH = D:\Julia\Julia-1.10.0\packages
  JULIA_PKG_SERVER = https://mirrors.pku.edu.cn/julia/
  JULIA_EDITOR = code
  JULIA_NUM_THREADS = 0

Additional context

But when I modified the version of DataInterpolations.jl to v5.0.0, the above code ran successfully.

([15.559944125537207, -4.53336304692014], [122.68735462923269, -195.97996798273402, 82.50456157341432, -223.49493920467143])

Additionally, when itp is defined outside the loss function, the above issue does not occur in either version. However, sometimes I need to define itp inside the loss function.

[ChrisRackauckas]

I believe @SouthEndMusic is handling this in #291

[sathvikbhagavan]

I think this works now:

julia> du0, dp = Zygote.gradient(loss, u0, p)
([15.796940059345747, -4.727910327153384], [124.23159522772528, -191.5089814958977, 84.94094858539026, -228.43161333505992])

(jl_D7NIX6) pkg> st
Status `/private/var/folders/4d/0gnbkbdj2hl6xyfx573_zwl80000gn/T/jl_D7NIX6/Project.toml`
  [82cc6244] DataInterpolations v6.5.0
  [1dea7af3] OrdinaryDiffEq v6.89.0
  [1ed8b502] SciMLSensitivity v7.68.0
  [e88e6eb3] Zygote v0.6.72

We can close this.