Example for integro-differential equation with NeuralPDE does not run

I am trying to use NeuralPDE for solving an integro-differential equation with the aim of later embedding a neural network in it. To begin with, I tried the example I found in the documentation, but it gives me an error:

using NeuralPDE, Flux, ModelingToolkit, Optimization, OptimizationOptimJL, DomainSets
import ModelingToolkit: Interval, infimum, supremum

@parameters t
@variables i(..)
Di = Differential(t)
Ii = Integral(t in DomainSets.ClosedInterval(0, t))
eq = Di(i(t)) + 2 * i(t) + 5 * Ii(i(t)) ~ 1
bcs = [i(0.0) ~ 0.0]
domains = [t ∈ Interval(0.0, 2.0)]
chain = Chain(Dense(1, 15, Flux.σ), Dense(15, 1)) |> f64

strategy_ = GridTraining(0.05)
discretization = PhysicsInformedNN(chain,
                                   strategy_)
@named pde_system = PDESystem(eq, bcs, domains, [t], [i(t)])
prob = NeuralPDE.discretize(pde_system, discretization)
callback = function (p, l)
    println("Current loss is: $l")
    return false
end
res = Optimization.solve(prob, BFGS(); callback = callback, maxiters = 100)
DimensionMismatch: array with ndims(x) == 2 >  0 cannot have dx::Number

Stacktrace:
  [1] (::Cubature.var"#17#18"{Bool, Bool, Int64, Float64, Float64, Int64, Int32, Ptr{Nothing}, Cubature.IntegrandData{IntegralsCubature.var"#7#15"{IntegralProblem{false, Vector{Float64}, IntegralsZygoteExt.var"#14#23"{Float64, Integrals.IntegralCache{false, NeuralPDE.var"#integrand_#286"{Vector{Float64}, NeuralPDE.var"#12#13", NeuralPDE.Phi{Optimisers.Restructure{Chain{Tuple{Dense{typeof(σ), Matrix{Float64}, Vector{Float64}}, Dense{typeof(identity), Matrix{Float64}, Vector{Float64}}}}, NamedTuple{(:layers,), Tuple{Tuple{NamedTuple{(:weight, :bias, :σ), Tuple{Int64, Int64, Tuple{}}}, NamedTuple{(:weight, :bias, :σ), Tuple{Int64, Int64, Tuple{}}}}}}}, Nothing}, Vector{Int64}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:cord, Symbol("##θ#293"), :phi, :derivative, :integral, :u, :p), NeuralPDE.var"#_RGF_ModTag", NeuralPDE.var"#_RGF_ModTag", (0xd0887b84, 0x430aacf6, 0x16011867, 0xdd06f795, 0xbb666625), Expr}, typeof(NeuralPDE.numeric_derivative)}, Vector{Float64}, Vector{Float64}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, IntegralsCubature.CubatureJLh, Integrals.ReCallVJP{Integrals.ZygoteVJP}, Base.Pairs{Symbol, Float64, Tuple{Symbol, Symbol}, NamedTuple{(:reltol, :abstol), Tuple{Float64, Float64}}}, Nothing}}, Vector{Float64}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, Vector{Float64}}}, Vector{Float64}, Vector{Float64}, Vector{Float64}, Vector{Float64}, Int64})()
    @ Cubature ~/.julia/packages/Cubature/5zwuu/src/Cubature.jl:215
[2] disable_sigint
    @ ./c.jl:473 [inlined]
  [3] cubature(xscalar::Bool, fscalar::Bool, vectorized::Bool, padaptive::Bool, fdim::Int64, f::IntegralsCubature.var"#7#15"{IntegralProblem{false, Vector{Float64}, IntegralsZygoteExt.var"#14#23"{Float64, Integrals.IntegralCache{false, NeuralPDE.var"#integrand_#286"{Vector{Float64}, NeuralPDE.var"#12#13", NeuralPDE.Phi{Optimisers.Restructure{Chain{Tuple{Dense{typeof(σ), Matrix{Float64}, Vector{Float64}}, Dense{typeof(identity), Matrix{Float64}, Vector{Float64}}}}, NamedTuple{(:layers,), Tuple{Tuple{NamedTuple{(:weight, :bias, :σ), Tuple{Int64, Int64, Tuple{}}}, NamedTuple{(:weight, :bias, :σ), Tuple{Int64, Int64, Tuple{}}}}}}}, Nothing}, Vector{Int64}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:cord, Symbol("##θ#293"), :phi, :derivative, :integral, :u, :p), NeuralPDE.var"#_RGF_ModTag", NeuralPDE.var"#_RGF_ModTag", (0xd0887b84, 0x430aacf6, 0x16011867, 0xdd06f795, 0xbb666625), Expr}, typeof(NeuralPDE.numeric_derivative)}, Vector{Float64}, Vector{Float64}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, IntegralsCubature.CubatureJLh, Integrals.ReCallVJP{Integrals.ZygoteVJP}, Base.Pairs{Symbol, Float64, Tuple{Symbol, Symbol}, NamedTuple{(:reltol, :abstol), Tuple{Float64, Float64}}}, Nothing}}, Vector{Float64}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, Vector{Float64}}, xmin_::Vector{Float64}, xmax_::Vector{Float64}, reqRelError::Float64, reqAbsError::Float64, maxEval::Int64, error_norm::Int32)
    @ Cubature ~/.julia/packages/Cubature/5zwuu/src/Cubature.jl:169
  [4] #hcubature#21
    @ ~/.julia/packages/Cubature/5zwuu/src/Cubature.jl:227 [inlined]
  [5] hcubature
    @ ~/.julia/packages/Cubature/5zwuu/src/Cubature.jl:227 [inlined]
  [6] __solvebp_call(prob::IntegralProblem{false, Vector{Float64}, IntegralsZygoteExt.var"#14#23"{Float64, Integrals.IntegralCache{false, NeuralPDE.var"#integrand_#286"{Vector{Float64}, NeuralPDE.var"#12#13", NeuralPDE.Phi{Optimisers.Restructure{Chain{Tuple{Dense{typeof(σ), Matrix{Float64}, Vector{Float64}}, Dense{typeof(identity), Matrix{Float64}, Vector{Float64}}}}, NamedTuple{(:layers,), Tuple{Tuple{NamedTuple{(:weight, :bias, :σ), Tuple{Int64, Int64, Tuple{}}}, NamedTuple{(:weight, :bias, :σ), Tuple{Int64, Int64, Tuple{}}}}}}}, Nothing}, Vector{Int64}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:cord, Symbol("##θ#293"), :phi, :derivative, :integral, :u, :p), NeuralPDE.var"#_RGF_ModTag", NeuralPDE.var"#_RGF_ModTag", (0xd0887b84, 0x430aacf6, 0x16011867, 0xdd06f795, 0xbb666625), Expr}, typeof(NeuralPDE.numeric_derivative)}, Vector{Float64}, Vector{Float64}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, IntegralsCubature.CubatureJLh, Integrals.ReCallVJP{Integrals.ZygoteVJP}, Base.Pairs{Symbol, Float64, Tuple{Symbol, Symbol}, NamedTuple{(:reltol, :abstol), Tuple{Float64, Float64}}}, Nothing}}, Vector{Float64}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, alg::IntegralsCubature.CubatureJLh, sensealg::Integrals.ReCallVJP{Integrals.ZygoteVJP}, lb::Vector{Float64}, ub::Vector{Float64}, p::Vector{Float64}; reltol::Float64, abstol::Float64, maxiters::Int64)
    @ IntegralsCubature ~/.julia/packages/IntegralsCubature/OPTxp/src/IntegralsCubature.jl:138
  [7] __solvebp_call
    @ ~/.julia/packages/IntegralsCubature/OPTxp/src/IntegralsCubature.jl:50 [inlined]
  [8] #__solvebp_call#4
    @ ~/.julia/packages/Integrals/d3rQd/src/common.jl:95 [inlined]
  [9] (::IntegralsZygoteExt.var"#quadrature_adjoint#16"{Base.Pairs{Symbol, Float64, Tuple{Symbol, Symbol}, NamedTuple{(:reltol, :abstol), Tuple{Float64, Float64}}}, Integrals.IntegralCache{false, NeuralPDE.var"#integrand_#286"{Vector{Float64}, NeuralPDE.var"#12#13", NeuralPDE.Phi{Optimisers.Restructure{Chain{Tuple{Dense{typeof(σ), Matrix{Float64}, Vector{Float64}}, Dense{typeof(identity), Matrix{Float64}, Vector{Float64}}}}, NamedTuple{(:layers,), Tuple{Tuple{NamedTuple{(:weight, :bias, :σ), Tuple{Int64, Int64, Tuple{}}}, NamedTuple{(:weight, :bias, :σ), Tuple{Int64, Int64, Tuple{}}}}}}}, Nothing}, Vector{Int64}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:cord, Symbol("##θ#293"), :phi, :derivative, :integral, :u, :p), NeuralPDE.var"#_RGF_ModTag", NeuralPDE.var"#_RGF_ModTag", (0xd0887b84, 0x430aacf6, 0x16011867, 0xdd06f795, 0xbb666625), Expr}, typeof(NeuralPDE.numeric_derivative)}, Vector{Float64}, Vector{Float64}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, IntegralsCubature.CubatureJLh, Integrals.ReCallVJP{Integrals.ZygoteVJP}, Base.Pairs{Symbol, Float64, Tuple{Symbol, Symbol}, NamedTuple{(:reltol, :abstol), Tuple{Float64, Float64}}}, Nothing}, IntegralsCubature.CubatureJLh, Integrals.ReCallVJP{Integrals.ZygoteVJP}, Vector{Float64}, Vector{Float64}, Vector{Float64}})(Δ::Array{Float64, 0})
    @ IntegralsZygoteExt ~/.julia/packages/Integrals/d3rQd/ext/IntegralsZygoteExt.jl:82

(full stacktrace is so long I could not paste it here)

My side question would be wether it makes sense to use NeuralPDE for automated discovery of a missing term in an integro-differential equation

@SathvikBhagavan did this get handled when we did the push for MTK v9 and fixed all of the tutorials? I believe it’s a test case?

The tutorial uses QuadratureTraining and it works Solving Integro Differential Equations · NeuralPDE.jl. I am not sure why it fails with GridTraining though. I will look into it.