fix: fix type instability in late_binding_update_u0_p

Author: AayushSabharwal

src/remake.jl

@@ -1083,6 +1083,10 @@ calling `SymbolicIndexingInterface.symbolic_container`, provided for dispatch. R
 the updated `newu0` and `newp`.
 """
 function late_binding_update_u0_p(prob, root_indp, u0, p, t0, newu0, newp)
+    if hasmethod(symbolic_container, Tuple{typeof(root_indp)}) &&
+       (sc = symbolic_container(root_indp)) !== root_indp
+        return late_binding_update_u0_p(prob.sc, u0, p, t0, newu0, newp)
+    end
     return newu0, newp
 end
 
@@ -1094,10 +1098,6 @@ Calls `late_binding_update_u0_p(prob, root_indp, u0, p, t0, newu0, newp)` after
 """
 function late_binding_update_u0_p(prob, u0, p, t0, newu0, newp)
     root_indp = prob
-    while hasmethod(symbolic_container, Tuple{typeof(root_indp)}) &&
-        (sc = symbolic_container(root_indp)) !== root_indp
-        root_indp = sc
-    end
     return late_binding_update_u0_p(prob, root_indp, u0, p, t0, newu0, newp)
 end