fix: fix type promotion in late_binding_update_u0_p

Author: AayushSabharwal

src/systems/nonlinear/initializesystem.jl

@@ -589,13 +589,31 @@ function SciMLBase.remake_initialization_data(
     return SciMLBase.remake_initialization_data(sys, kws, newu0, t0, newp, newu0, newp)
 end
 
+function promote_u0_p(u0, p::MTKParameters, t0)
+    u0 = DiffEqBase.promote_u0(u0, p.tunable, t0)
+    u0 = DiffEqBase.promote_u0(u0, p.initials, t0)
+
+    tunables = DiffEqBase.promote_u0(p.tunable, u0, t0)
+    initials = DiffEqBase.promote_u0(p.initials, u0, t0)
+    p = SciMLStructures.replace(SciMLStructures.Tunable(), p, tunables)
+    p = SciMLStructures.replace(SciMLStructures.Initials(), p, initials)
+
+    return u0, p
+end
+
+function promote_u0_p(u0, p::AbstractArray, t0)
+    return DiffEqBase.promote_u0(u0, p, t0), DiffEqBase.promote_u0(p, u0, t0)
+end
+
 function SciMLBase.late_binding_update_u0_p(
         prob, sys::AbstractSystem, u0, p, t0, newu0, newp)
     supports_initialization(sys) || return newu0, newp
 
     initdata = prob.f.initialization_data
     meta = initdata === nothing ? nothing : initdata.metadata
 
+    newu0, newp = promote_u0_p(newu0, newp, t0)
+
     # non-symbolic u0 updates initials...
     if !(eltype(u0) <: Pair)
         # if `p` is not provided or is symbolic
@@ -605,12 +623,7 @@ function SciMLBase.late_binding_update_u0_p(
         meta = initdata.metadata
         meta isa InitializationMetadata || return newu0, newp
         newp = p === missing ? copy(newp) : newp
-        initials, repack, alias = SciMLStructures.canonicalize(
-            SciMLStructures.Initials(), newp)
-        if eltype(initials) != eltype(newu0)
-            initials = DiffEqBase.promote_u0(initials, newu0, t0)
-            newp = repack(initials)
-        end
+
         if length(newu0) != length(prob.u0)
             throw(ArgumentError("Expected `newu0` to be of same length as unknowns ($(length(prob.u0))). Got $(typeof(newu0)) of length $(length(newu0))"))
         end
@@ -619,17 +632,6 @@ function SciMLBase.late_binding_update_u0_p(
     end
 
     newp = p === missing ? copy(newp) : newp
-    newu0 = DiffEqBase.promote_u0(newu0, newp, t0)
-    tunables, repack, alias = SciMLStructures.canonicalize(SciMLStructures.Tunable(), newp)
-    if eltype(tunables) != eltype(newu0)
-        tunables = DiffEqBase.promote_u0(tunables, newu0, t0)
-        newp = repack(tunables)
-    end
-    initials, repack, alias = SciMLStructures.canonicalize(SciMLStructures.Initials(), newp)
-    if eltype(initials) != eltype(newu0)
-        initials = DiffEqBase.promote_u0(initials, newu0, t0)
-        newp = repack(initials)
-    end
 
     allsyms = all_symbols(sys)
     for (k, v) in u0