chore: treat prob as immutable

Author: DhairyaLGandhi

Project.toml

@@ -31,6 +31,7 @@ SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [weakdeps]
 ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
@@ -41,7 +42,6 @@ PartialFunctions = "570af359-4316-4cb7-8c74-252c00c2016b"
 PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"
 PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 RCall = "6f49c342-dc21-5d91-9882-a32aef131414"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [extensions]
 SciMLBaseChainRulesCoreExt = "ChainRulesCore"

ext/SciMLBaseChainRulesCoreExt.jl

@@ -58,6 +58,7 @@ end
 
 function ChainRulesCore.rrule(::Type{ODEProblem}, args...; kwargs...)
     function ODEProblemAdjoint(ȳ)
+        @show "some con"
         (NoTangent(), ȳ.f, ȳ.u0, ȳ.tspan, ȳ.p, ȳ.kwargs, ȳ.problem_type)
     end
 

ext/SciMLBaseZygoteExt.jl

@@ -300,7 +300,7 @@ end
     ∇responsible_map(__context__, f, args...)
 end
 
-@_adjoint_keepthunks function Zygote.literal_getfield(x::ODEProblem, ::Val{f}) where f
+@_adjoint_keepthunks function Zygote.literal_getfield(x::SciMLBase.AbstractSciMLProblem, ::Val{f}) where f
   val = getfield(x, f)
   function back(Δ)
     Zygote.accum_param(__context__, val, Δ) === nothing && return
@@ -311,7 +311,7 @@ end
     else
       dx = Zygote.grad_mut(__context__, x)
       dx[] = (; dx[]..., pair(Val(f), Zygote.accum(getfield(dx[], f), Δ))...)
-      return (dx,nothing)
+      return (dx[],nothing)
     end
   end
   Zygote.unwrap(val), back

src/SciMLBase.jl

@@ -44,6 +44,8 @@ import SciMLOperators:
 
 @reexport using SciMLOperators
 
+using Zygote
+
 function __solve end
 function __init end
 

src/initialization.jl

@@ -260,9 +260,8 @@ function get_initial_values(prob, valp, f, alg::OverrideInit,
     end
     nlsol, success = solve_initialization(initdata, initprob, alg; reltol, abstol, nlsolve_alg )
 
-    nlsol2 = prob.f.initialization_data.initializeprob
     if initdata.initializeprobmap !== nothing
-        u02 = initdata.initializeprobmap(nlsol2)
+        u02 = initdata.initializeprobmap(nlsol)
     end
     if initdata.initializeprobpmap !== nothing
         p2 = initdata.initializeprobpmap(valp, nlsol)

src/remake.jl

@@ -703,6 +703,8 @@ function remake(prob::NonlinearProblem;
     if problem_type === missing
         problem_type = prob.problem_type
     end
+    # error()
+    # @show f
 
     prob = if kwargs === missing
         NonlinearProblem{isinplace(prob)}(f = f, u0 = newu0, p = newp,
@@ -1206,18 +1208,25 @@ function maybe_eager_initialize_problem(prob::AbstractSciMLProblem, initializati
     if lazy_initialization === nothing
         lazy_initialization = !is_trivial_initialization(initialization_data)
     end
-    if initialization_data !== nothing && !lazy_initialization &&
-       (!is_time_dependent(prob) || current_time(prob) !== nothing)
+    cond = initialization_data !== nothing && !lazy_initialization &&
+            (!is_time_dependent(prob) || current_time(prob) !== nothing)
+    @show cond
+    if cond
+    #    @show "in maybe_eager_initialize_problem"
         u0, p, _ = get_initial_values(
             prob, prob, prob.f, OverrideInit(), Val(isinplace(prob)))
-        if u0 !== nothing && eltype(u0) == Any && isempty(u0)
-            u0 = nothing
-        end
+        # if u0 !== nothing && eltype(u0) == Any && isempty(u0)
+        #     u0 = nothing
+        # end
     else
-        u0 = state_values(prob)
-        p = parameter_values(prob)
+        u02 = state_values(prob)
+        p2 = parameter_values(prob)
     end
-    return u0, p
+    # @show p
+
+    u03 = cond ? u0 : u02
+    p3 = cond ? p : p2
+    return u03, p3
 end
 
 function remake(thing::AbstractJumpProblem; kwargs...)