Type inference failure for `remake(::ODEProblem; u0, p)` with MTK generated problems

[SebastianM-C]

Describe the bug 🐞

Similar to #981, this is the more advanced version where we actually have to change some values. The example below is contrived, but passing u0 and p with these types is encountered when using setsym_oop. While the u0 and p values are made type stable by SciML/ModelingToolkit.jl#3546 & SciML/SymbolicIndexingInterface.jl#118 (with something like SciML/SymbolicIndexingInterface.jl#119 on top), the issue is that the remake for the initialization data is very dynamic and completely loses the gains from the previous PRs.

The starting point is

meta = get_metadata(oldinitsys)
if meta isa InitializationSystemMetadata && meta.oop_reconstruct_u0_p !== nothing
    reconstruct_fn = meta.oop_reconstruct_u0_p
else
    reconstruct_fn = ReconstructInitializeprob(sys, oldinitsys)
end

from which we get a Union{Nothing, ReconstructInitializeprob} that when called with another Union in the ProblemState (due to history_fn) completely gives up on inference and we get new_initp::Any, which poisons the rest.

Expected behavior

remake should be inferred.

Minimal Reproducible Example 👇

Without MRE, we would only be able to help you to a limited extent, and attention to the issue would be limited. to know more about MRE refer to wikipedia and stackoverflow.

using SciMLBase

prob = ODEProblem((u, p, t) -> -u, [1.0], (0.0, 1.0))
@code_warntype remake(prob, u0=prob.u0, p=prob.p) # no issues

using ModelingToolkit
using ModelingToolkit: t_nounits as t, D_nounits as D

@parameters α=1 β=1 γ=1 δ=1
@variables x(t)=1 y(t)=1
eqs = [D(x) ~ α * x - β * x * y, D(y) ~ -δ * y + γ * x * y]
@named sys = ODESystem(eqs, t)
prob = ODEProblem(complete(sys), [], (0.,1))

@code_warntype remake(prob, u0=prob.u0, p=prob.p) # Any

Error & Stacktrace ⚠️

MethodInstance for Core.kwcall(::@NamedTuple{u0::Vector{Float64}, p::MTKParameters{Vector{Float64}, Vector{Float64}, Tuple{}, Tuple{}, Tuple{}, Tuple{}}}, ::typeof(remake), ::ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, MTKParameters{Vector{Float64}, Vector{Float64}, Tuple{}, Tuple{}, Tuple{}, Tuple{}}, ODEFunction{true, SciMLBase.AutoSpecialize, ModelingToolkit.GeneratedFunctionWrapper{(2, 3, true), RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:__mtk_arg_1, :___mtkparameters___, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0x42dc2497, 0x45642c41, 0x10ed7100, 0x37cb3c7b, 0x39fcf7e9), Nothing}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋out, :__mtk_arg_1, :___mtkparameters___, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0x00f7494c, 0x1d4dc422, 0xd4b8b833, 0xca7f69f5, 0x2249aa1c), Nothing}}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, ModelingToolkit.ObservedFunctionCache{ODESystem}, Nothing, ODESystem, SciMLBase.OverrideInitData{NonlinearProblem{Nothing, true, MTKParameters{Vector{Float64}, StaticArraysCore.SizedVector{0, Float64, Vector{Float64}}, Tuple{}, Tuple{}, Tuple{}, Tuple{}}, NonlinearFunction{true, SciMLBase.FullSpecialize, ModelingToolkit.GeneratedFunctionWrapper{(2, 2, true), RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:__mtk_arg_1, :___mtkparameters___), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0x85e41679, 0xe3baaf69, 0x36273733, 0xfa33874a, 0x6249a58a), Nothing}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋out, :__mtk_arg_1, :___mtkparameters___), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0xa427c4e8, 0x725f824c, 0xe49d1abc, 0x4f823061, 0x3d8f5fff), Nothing}}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, ModelingToolkit.ObservedFunctionCache{NonlinearSystem}, Nothing, NonlinearSystem, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, ModelingToolkit.InitializationSystemMetadata{ModelingToolkit.ReconstructInitializeprob{SymbolicIndexingInterface.MultipleGetters{ContinuousTimeseries, Vector{SymbolicIndexingInterface.AbstractGetIndexer}}, SymbolicIndexingInterface.OOPSetter{false, NonlinearSystem, Vector{ModelingToolkit.ParameterIndex{SciMLStructures.Tunable, B, Int64} where B}}}}}, ModelingToolkit.UpdateInitializeprob{SymbolicIndexingInterface.MultipleGetters{ContinuousTimeseries, Vector{SymbolicIndexingInterface.AbstractGetIndexer}}, SymbolicIndexingInterface.MultipleSetters{Vector{SymbolicIndexingInterface.ParameterHookWrapper{S, SymbolicUtils.BasicSymbolic{Real}} where S}}}, ComposedFunction{typeof(identity), SymbolicIndexingInterface.TimeIndependentObservedFunction{ModelingToolkit.GeneratedFunctionWrapper{(2, 2, true), RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:__mtk_arg_1, :___mtkparameters___), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0xb1e81c3a, 0x1774638f, 0xdf6809cd, 0xabae0e03, 0x7fcc3406), Nothing}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋out, :__mtk_arg_1, :___mtkparameters___), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0x11d787e2, 0x05a9fa01, 0xefd56ed6, 0xcbb4f113, 0xd37eeca2), Nothing}}}}, Nothing}, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardODEProblem})
  from kwcall(::NamedTuple, ::typeof(remake), prob::ODEProblem) @ SciMLBase ~/.julia/dev/SciMLBase/src/remake.jl:214
Arguments
  _::Core.Const(Core.kwcall)
  @_2::@NamedTuple{u0::Vector{Float64}, p::MTKParameters{Vector{Float64}, Vector{Float64}, Tuple{}, Tuple{}, Tuple{}, Tuple{}}}
  @_3::Core.Const(SciMLBase.remake)
  prob::ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, MTKParameters{Vector{Float64}, Vector{Float64}, Tuple{}, Tuple{}, Tuple{}, Tuple{}}, ODEFunction{true, SciMLBase.AutoSpecialize, ModelingToolkit.GeneratedFunctionWrapper{(2, 3, true), RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:__mtk_arg_1, :___mtkparameters___, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0x42dc2497, 0x45642c41, 0x10ed7100, 0x37cb3c7b, 0x39fcf7e9), Nothing}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋out, :__mtk_arg_1, :___mtkparameters___, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0x00f7494c, 0x1d4dc422, 0xd4b8b833, 0xca7f69f5, 0x2249aa1c), Nothing}}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, ModelingToolkit.ObservedFunctionCache{ODESystem}, Nothing, ODESystem, SciMLBase.OverrideInitData{NonlinearProblem{Nothing, true, MTKParameters{Vector{Float64}, StaticArraysCore.SizedVector{0, Float64, Vector{Float64}}, Tuple{}, Tuple{}, Tuple{}, Tuple{}}, NonlinearFunction{true, SciMLBase.FullSpecialize, ModelingToolkit.GeneratedFunctionWrapper{(2, 2, true), RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:__mtk_arg_1, :___mtkparameters___), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0x85e41679, 0xe3baaf69, 0x36273733, 0xfa33874a, 0x6249a58a), Nothing}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋out, :__mtk_arg_1, :___mtkparameters___), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0xa427c4e8, 0x725f824c, 0xe49d1abc, 0x4f823061, 0x3d8f5fff), Nothing}}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, ModelingToolkit.ObservedFunctionCache{NonlinearSystem}, Nothing, NonlinearSystem, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, ModelingToolkit.InitializationSystemMetadata{ModelingToolkit.ReconstructInitializeprob{SymbolicIndexingInterface.MultipleGetters{ContinuousTimeseries, Vector{SymbolicIndexingInterface.AbstractGetIndexer}}, SymbolicIndexingInterface.OOPSetter{false, NonlinearSystem, Vector{ModelingToolkit.ParameterIndex{SciMLStructures.Tunable, B, Int64} where B}}}}}, ModelingToolkit.UpdateInitializeprob{SymbolicIndexingInterface.MultipleGetters{ContinuousTimeseries, Vector{SymbolicIndexingInterface.AbstractGetIndexer}}, SymbolicIndexingInterface.MultipleSetters{Vector{SymbolicIndexingInterface.ParameterHookWrapper{S, SymbolicUtils.BasicSymbolic{Real}} where S}}}, ComposedFunction{typeof(identity), SymbolicIndexingInterface.TimeIndependentObservedFunction{ModelingToolkit.GeneratedFunctionWrapper{(2, 2, true), RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:__mtk_arg_1, :___mtkparameters___), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0xb1e81c3a, 0x1774638f, 0xdf6809cd, 0xabae0e03, 0x7fcc3406), Nothing}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋out, :__mtk_arg_1, :___mtkparameters___), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0x11d787e2, 0x05a9fa01, 0xefd56ed6, 0xcbb4f113, 0xd37eeca2), Nothing}}}}, Nothing}, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardODEProblem}
Locals
  f::Union{}
  u0::Union{}
  tspan::Union{}
  p::Union{}
  kwargs::Union{}
  interpret_symbolicmap::Union{}
  build_initializeprob::Union{}
  use_defaults::Union{}
  lazy_initialization::Union{}
  _kwargs...::Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}
  @_15::Union{Missing, Nothing, Bool, Type{Val{true}}, MTKParameters{Vector{Float64}, Vector{Float64}, Tuple{}, Tuple{}, Tuple{}, Tuple{}}, Vector{Float64}}
Body::Any
1 ──       Core.NewvarNode(:(f))
│          Core.NewvarNode(:(u0))
│          Core.NewvarNode(:(tspan))
│          Core.NewvarNode(:(p))
│          Core.NewvarNode(:(kwargs))
│          Core.NewvarNode(:(interpret_symbolicmap))
│          Core.NewvarNode(:(build_initializeprob))
│          Core.NewvarNode(:(use_defaults))
│          Core.NewvarNode(:(lazy_initialization))
│          Core.NewvarNode(:(_kwargs...))
│          Core.NewvarNode(:(@_15))
│    %12 = Core.isdefined(@_2, :f)::Core.Const(false)
└───       goto #3 if not %12
2 ──       Core.Const(:(@_15 = Core.getfield(@_2, :f)))
└───       Core.Const(:(goto %17))
3 ┄─       (@_15 = SciMLBase.missing)
│    %17 = @_15::Core.Const(missing)
│    %18 = Core.isdefined(@_2, :u0)::Core.Const(true)
└───       goto #5 if not %18
4 ──       (@_15 = Core.getfield(@_2, :u0))
└───       goto #6
5 ──       Core.Const(:(@_15 = SciMLBase.missing))
6 ┄─ %23 = @_15::Vector{Float64}
│    %24 = Core.isdefined(@_2, :tspan)::Core.Const(false)
└───       goto #8 if not %24
7 ──       Core.Const(:(@_15 = Core.getfield(@_2, :tspan)))
└───       Core.Const(:(goto %29))
8 ┄─       (@_15 = SciMLBase.missing)
│    %29 = @_15::Core.Const(missing)
│    %30 = Core.isdefined(@_2, :p)::Core.Const(true)
└───       goto #10 if not %30
9 ──       (@_15 = Core.getfield(@_2, :p))
└───       goto #11
10 ─       Core.Const(:(@_15 = SciMLBase.missing))
11 ┄ %35 = @_15::MTKParameters{Vector{Float64}, Vector{Float64}, Tuple{}, Tuple{}, Tuple{}, Tuple{}}
│    %36 = Core.isdefined(@_2, :kwargs)::Core.Const(false)
└───       goto #13 if not %36
12 ─       Core.Const(:(@_15 = Core.getfield(@_2, :kwargs)))
└───       Core.Const(:(goto %41))
13 ┄       (@_15 = SciMLBase.missing)
│    %41 = @_15::Core.Const(missing)
│    %42 = Core.isdefined(@_2, :interpret_symbolicmap)::Core.Const(false)
└───       goto #15 if not %42
14 ─       Core.Const(:(@_15 = Core.getfield(@_2, :interpret_symbolicmap)))
└───       Core.Const(:(goto %47))
15 ┄       (@_15 = true)
│    %47 = @_15::Core.Const(true)
│    %48 = Core.isdefined(@_2, :build_initializeprob)::Core.Const(false)
└───       goto #17 if not %48
16 ─       Core.Const(:(@_15 = Core.getfield(@_2, :build_initializeprob)))
└───       Core.Const(:(goto %53))
17 ┄       (@_15 = Core.apply_type(SciMLBase.Val, true))
│    %53 = @_15::Core.Const(Val{true})
│    %54 = Core.isdefined(@_2, :use_defaults)::Core.Const(false)
└───       goto #19 if not %54
18 ─       Core.Const(:(@_15 = Core.getfield(@_2, :use_defaults)))
└───       Core.Const(:(goto %59))
19 ┄       (@_15 = false)
│    %59 = @_15::Core.Const(false)
│    %60 = Core.isdefined(@_2, :lazy_initialization)::Core.Const(false)
└───       goto #21 if not %60
20 ─       Core.Const(:(@_15 = Core.getfield(@_2, :lazy_initialization)))
└───       Core.Const(:(goto %65))
21 ┄       (@_15 = SciMLBase.nothing)
│    %65 = @_15::Core.Const(nothing)
│    %66 = (:f, :u0, :tspan, :p, :kwargs, :interpret_symbolicmap, :build_initializeprob, :use_defaults, :lazy_initialization)::Core.Const((:f, :u0, :tspan, :p, :kwargs, :interpret_symbolicmap, :build_initializeprob, :use_defaults, :lazy_initialization))
│    %67 = Core.apply_type(Core.NamedTuple, %66)::Core.Const(NamedTuple{(:f, :u0, :tspan, :p, :kwargs, :interpret_symbolicmap, :build_initializeprob, :use_defaults, :lazy_initialization)})
│    %68 = Base.structdiff(@_2, %67)::Core.Const(NamedTuple())
│          (_kwargs... = Base.pairs(%68))
│    %70 = SciMLBase.:(var"#remake#735")::Core.Const(SciMLBase.var"#remake#735")
│    %71 = _kwargs...::Core.Const(Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}())
│    %72 = (%70)(%17, %23, %29, %35, %41, %47, %53, %59, %65, %71, @_3, prob)::Any
└───       return %72

The relevant part from @descend would be

Environment (please complete the following information):

• Output of using Pkg; Pkg.status()

• Output of using Pkg; Pkg.status(; mode = PKGMODE_MANIFEST)

• Output of versioninfo()

Additional context

Add any other context about the problem here.

[SebastianM-C]

This seems fixed now.