@@ -5,65 +5,61 @@ using DiffEqBase: DiffEqBase
55using ReverseDiff: ReverseDiff, TrackedArray, TrackedReal
66using SciMLBase: ReverseDiffOriginator, NonlinearProblem, NonlinearLeastSquaresProblem
77using SimpleNonlinearSolve: SimpleNonlinearSolve
8+ import SimpleNonlinearSolve: __internal_solve_up
89
9- function SimpleNonlinearSolve . __internal_solve_up (
10+ function __internal_solve_up (
1011 prob:: Union{NonlinearProblem, NonlinearLeastSquaresProblem} , sensealg,
1112 u0:: TrackedArray , u0_changed, p:: TrackedArray , p_changed, alg, args... ; kwargs... )
12- return ReverseDiff. track (SimpleNonlinearSolve . __internal_solve_up, prob, sensealg,
13- u0, u0_changed, p, p_changed, alg, args... ; kwargs... )
13+ return ReverseDiff. track (__internal_solve_up, prob, sensealg, u0 ,
14+ u0_changed, p, p_changed, alg, args... ; kwargs... )
1415end
1516
16- function SimpleNonlinearSolve . __internal_solve_up (
17+ function __internal_solve_up (
1718 prob:: Union{NonlinearProblem, NonlinearLeastSquaresProblem} , sensealg,
1819 u0, u0_changed, p:: TrackedArray , p_changed, alg, args... ; kwargs... )
19- return ReverseDiff. track (SimpleNonlinearSolve . __internal_solve_up, prob, sensealg,
20- u0, u0_changed, p, p_changed, alg, args... ; kwargs... )
20+ return ReverseDiff. track (__internal_solve_up, prob, sensealg, u0 ,
21+ u0_changed, p, p_changed, alg, args... ; kwargs... )
2122end
2223
23- function SimpleNonlinearSolve . __internal_solve_up (
24+ function __internal_solve_up (
2425 prob:: Union{NonlinearProblem, NonlinearLeastSquaresProblem} , sensealg,
2526 u0:: TrackedArray , u0_changed, p, p_changed, alg, args... ; kwargs... )
26- return ReverseDiff. track (SimpleNonlinearSolve . __internal_solve_up, prob, sensealg,
27- u0, u0_changed, p, p_changed, alg, args... ; kwargs... )
27+ return ReverseDiff. track (__internal_solve_up, prob, sensealg, u0 ,
28+ u0_changed, p, p_changed, alg, args... ; kwargs... )
2829end
2930
30- function SimpleNonlinearSolve. __internal_solve_up (
31- prob:: Union{NonlinearProblem, NonlinearLeastSquaresProblem} ,
31+ function __internal_solve_up (prob:: Union{NonlinearProblem, NonlinearLeastSquaresProblem} ,
3232 sensealg, u0:: AbstractArray{<:TrackedReal} , u0_changed,
3333 p:: AbstractArray{<:TrackedReal} , p_changed, alg, args... ; kwargs... )
34- return SimpleNonlinearSolve. __internal_solve_up (
35- prob, sensealg, ArrayInterface. aos_to_soa (u0), true ,
34+ return __internal_solve_up (prob, sensealg, ArrayInterface. aos_to_soa (u0), true ,
3635 ArrayInterface. aos_to_soa (p), true , alg, args... ; kwargs... )
3736end
3837
39- function SimpleNonlinearSolve . __internal_solve_up (
38+ function __internal_solve_up (
4039 prob:: Union{NonlinearProblem, NonlinearLeastSquaresProblem} , sensealg, u0,
4140 u0_changed, p:: AbstractArray{<:TrackedReal} , p_changed, alg, args... ; kwargs... )
42- return SimpleNonlinearSolve. __internal_solve_up (
43- prob, sensealg, u0, true , ArrayInterface. aos_to_soa (p),
41+ return __internal_solve_up (prob, sensealg, u0, true , ArrayInterface. aos_to_soa (p),
4442 true , alg, args... ; kwargs... )
4543end
4644
47- function SimpleNonlinearSolve. __internal_solve_up (
48- prob:: Union{NonlinearProblem, NonlinearLeastSquaresProblem} ,
45+ function __internal_solve_up (prob:: Union{NonlinearProblem, NonlinearLeastSquaresProblem} ,
4946 sensealg, u0:: AbstractArray{<:TrackedReal} ,
5047 u0_changed, p, p_changed, alg, args... ; kwargs... )
51- return SimpleNonlinearSolve. __internal_solve_up (
52- prob, sensealg, u0, true , ArrayInterface. aos_to_soa (p),
48+ return __internal_solve_up (prob, sensealg, u0, true , ArrayInterface. aos_to_soa (p),
5349 true , alg, args... ; kwargs... )
5450end
5551
56- ReverseDiff. @grad function SimpleNonlinearSolve . __internal_solve_up (
52+ ReverseDiff. @grad function __internal_solve_up (
5753 prob:: Union{NonlinearProblem, NonlinearLeastSquaresProblem} ,
5854 sensealg, u0, u0_changed, p, p_changed, alg, args... ; kwargs... )
5955 out, ∇internal = DiffEqBase. _solve_adjoint (
6056 prob, sensealg, ReverseDiff. value (u0), ReverseDiff. value (p),
6157 ReverseDiffOriginator (), alg, args... ; kwargs... )
62- function ∇SimpleNonlinearSolve . __internal_solve_up (_args... )
58+ function ∇__internal_solve_up (_args... )
6359 ∂prob, ∂sensealg, ∂u0, ∂p, ∂originator, ∂args... = ∇internal (_args... )
6460 return (∂prob, ∂sensealg, ∂u0, nothing , ∂p, nothing , nothing , ∂args... )
6561 end
66- return Array (out), ∇SimpleNonlinearSolve . __internal_solve_up
62+ return Array (out), ∇__internal_solve_up
6763end
6864
6965end
0 commit comments