Don't specialize on original and stats for polyalg

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "NonlinearSolve"
 uuid = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 authors = ["SciML"]
-version = "3.8.2"
+version = "3.8.3"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

src/default.jl

@@ -156,7 +156,7 @@ end
                             $(u_result_syms[i]) = $(sol_syms[i]).u
                         end
                         fu = get_fu($(cache_syms[i]))
-                        return SciMLBase.build_solution(
+                        return __build_solution_less_specialize(
                             $(sol_syms[i]).prob, cache.alg, $(u_result_syms[i]),
                             fu; retcode = $(sol_syms[i]).retcode, stats,
                             original = $(sol_syms[i]), trace = $(sol_syms[i]).trace)
@@ -196,7 +196,8 @@ end
                 copyto!(cache.u0, u)
                 u = cache.u0
             end
-            return SciMLBase.build_solution(cache.caches[idx].prob, cache.alg, u, fus[idx];
+            return __build_solution_less_specialize(
+                cache.caches[idx].prob, cache.alg, u, fus[idx];
                 retcode, stats, cache.caches[idx].trace)
         end)
 
@@ -283,7 +284,7 @@ for (probType, pType) in ((:NonlinearProblem, :NLS), (:NonlinearLeastSquaresProb
                                 else
                                     $(u_result_syms[i]) = $(cur_sol).u
                                 end
-                                return SciMLBase.build_solution(
+                                return __build_solution_less_specialize(
                                     prob, alg, $(u_result_syms[i]), $(cur_sol).resid;
                                     $(cur_sol).retcode, $(cur_sol).stats,
                                     original = $(cur_sol), trace = $(cur_sol).trace)
@@ -316,7 +317,7 @@ for (probType, pType) in ((:NonlinearProblem, :NLS), (:NonlinearLeastSquaresProb
                             else
                                 $(u_result_syms[i]) = $(sol_syms[i]).u
                             end
-                            return SciMLBase.build_solution(
+                            return __build_solution_less_specialize(
                                 prob, alg, $(u_result_syms[i]), $(sol_syms[i]).resid;
                                 $(sol_syms[i]).retcode, $(sol_syms[i]).stats,
                                 $(sol_syms[i]).trace, original = $(sol_syms[i]))

src/utils.jl

@@ -176,3 +176,16 @@ Determine the chunk size for ForwardDiff and PolyesterForwardDiff based on the i
 """
 @inline pickchunksize(x) = pickchunksize(length(x))
 @inline pickchunksize(x::Int) = ForwardDiff.pickchunksize(x)
+
+# Original is often determined on runtime information especially for PolyAlgorithms so it
+# is best to never specialize on that
+function __build_solution_less_specialize(prob::AbstractNonlinearProblem, alg, u, resid;
+        retcode = ReturnCode.Default, original = nothing, left = nothing,
+        right = nothing, stats = nothing, trace = nothing, kwargs...)
+    T = eltype(eltype(u))
+    N = ndims(u)
+
+    return SciMLBase.NonlinearSolution{T, N, typeof(u), typeof(resid), typeof(prob),
+        typeof(alg), Any, typeof(left), Any, typeof(trace)}(
+        u, resid, prob, alg, retcode, original, left, right, stats, trace)
+end

test/misc/polyalg_tests.jl

@@ -224,6 +224,7 @@ end
 
     @test all(!isnan, sol.u)
     @test !SciMLBase.successful_retcode(sol.retcode)
+    @inferred solve(prob)
 end
 
 @testitem "[OOP] Infeasible" setup=[InfeasibleFunction] begin
@@ -235,6 +236,7 @@ end
 
     @test all(!isnan, sol.u)
     @test !SciMLBase.successful_retcode(sol.retcode)
+    @inferred solve(prob)
 
     u0 = @SVector [0.0, 0.0, 0.0]
     prob = NonlinearProblem(f1_infeasible, u0)
@@ -243,6 +245,7 @@ end
         sol = solve(prob)
         @test all(!isnan, sol.u)
         @test !SciMLBase.successful_retcode(sol.retcode)
+        @inferred solve(prob)
     catch err
         @test err isa LinearAlgebra.SingularException
     end