fix caching issues

Author: jClugstor

debug_dual.jl

@@ -0,0 +1,29 @@
+using LinearSolve
+using ForwardDiff
+using Test
+
+function h(p)
+    (A = [p[1] p[2]+1 p[2]^3;
+          3*p[1] p[1]+5 p[2] * p[1]-4;
+          p[2]^2 9*p[1] p[2]],
+        b = [p[1] + 1, p[2] * 2, p[1]^2])
+end
+
+A, b = h([ForwardDiff.Dual(5.0, 1.0, 0.0), ForwardDiff.Dual(5.0, 0.0, 1.0)])
+
+println("A:")
+display(A)
+println("\nb:")
+display(b)
+
+prob = LinearProblem(A, b)
+overload_x_p = solve(prob)
+backslash_x_p = A \ b
+
+println("\nExpected result (backslash):")
+display(backslash_x_p)
+println("\nLinearSolve result:")
+display(overload_x_p)
+
+println("\nDifference:")
+display(overload_x_p - backslash_x_p)

ext/LinearSolveForwardDiffExt.jl

@@ -48,6 +48,7 @@ LinearSolve.@concrete mutable struct DualLinearCache{DT}
     # Cached intermediate values for calculations
     rhs_list
     dual_u0_cache
+    primal_u_cache
     primal_b_cache
 
     dual_A
@@ -58,19 +59,19 @@ end
 function linearsolve_forwarddiff_solve!(cache::DualLinearCache, alg, args...; kwargs...)
     # Solve the primal problem
     cache.dual_u0_cache .= cache.linear_cache.u
-    sol = solve!(cache.linear_cache, alg, args...; kwargs...)  
+    sol = solve!(cache.linear_cache, alg, args...; kwargs...)
 
     cache.primal_b_cache .= cache.linear_cache.b
-    uu = sol.u
+    cache.primal_u_cache .= cache.linear_cache.u
 
     # Store solution metadata without copying - we'll return this
     primal_sol = sol
 
-    # Solves Dual partials separately 
+    # Solves Dual partials separately
     ∂_A = cache.partials_A
     ∂_b = cache.partials_b
 
-    xp_linsolve_rhs!(uu, ∂_A, ∂_b, cache)
+    xp_linsolve_rhs!(∂_A, ∂_b, cache)
 
     rhs_list = cache.rhs_list
 
@@ -83,11 +84,13 @@ function linearsolve_forwarddiff_solve!(cache::DualLinearCache, alg, args...; kw
 
     # Reset to the original `b` and `u`, users will expect that `b` doesn't change if they don't tell it to
     cache.linear_cache.b .= cache.primal_b_cache
+    cache.linear_cache.u .= cache.primal_u_cache
 
     return primal_sol
 end
 
-function xp_linsolve_rhs!(uu, ∂_A::Union{<:Partials, <:AbstractArray{<:Partials}},
+function xp_linsolve_rhs!(
+        ∂_A::Union{<:Partials, <:AbstractArray{<:Partials}},
         ∂_b::Union{<:Partials, <:AbstractArray{<:Partials}}, cache::DualLinearCache)
 
     # Update cached partials lists
@@ -100,14 +103,14 @@ function xp_linsolve_rhs!(uu, ∂_A::Union{<:Partials, <:AbstractArray{<:Partial
     # Compute rhs = b - A*uu using five-argument mul!
     for i in eachindex(b_list)
         cache.rhs_list[i] .= b_list[i]
-        mul!(cache.rhs_list[i], A_list[i], uu, -1, 1)
+        mul!(cache.rhs_list[i], A_list[i], cache.primal_u_cache, -1, 1)
     end
 
     return cache.rhs_list
 end
 
 function xp_linsolve_rhs!(
-        uu, ∂_A::Union{<:Partials, <:AbstractArray{<:Partials}},
+        ∂_A::Union{<:Partials, <:AbstractArray{<:Partials}},
         ∂_b::Nothing, cache::DualLinearCache)
 
     # Update cached partials list for A
@@ -116,14 +119,14 @@ function xp_linsolve_rhs!(
 
     # Compute rhs = -A*uu using five-argument mul!
     for i in eachindex(A_list)
-        mul!(cache.rhs_list[i], A_list[i], uu, -1, 0)
+        mul!(cache.rhs_list[i], A_list[i], cache.primal_u_cache, -1, 0)
     end
 
     return cache.rhs_list
 end
 
 function xp_linsolve_rhs!(
-        uu, ∂_A::Nothing, ∂_b::Union{<:Partials, <:AbstractArray{<:Partials}},
+        ∂_A::Nothing, ∂_b::Union{<:Partials, <:AbstractArray{<:Partials}},
         cache::DualLinearCache)
 
     # Update cached partials list for b
@@ -219,7 +222,6 @@ function __dual_init(
     partials_b_list = !isnothing(∂_b) ? partials_to_list(∂_b) : nothing
 
     # Determine size and type for rhs_list
-    n_partials = 0
     if !isnothing(partials_A_list)
         n_partials = length(partials_A_list)
         rhs_list = [similar(non_partial_cache.b) for _ in 1:n_partials]
@@ -240,6 +242,7 @@ function __dual_init(
         rhs_list,
         similar(new_b),
         similar(new_b),
+        similar(new_b),
         A,
         b,
         zeros(dual_type, length(b))
@@ -254,7 +257,7 @@ function SciMLBase.solve!(
         cache::DualLinearCache{DT}, alg::SciMLLinearSolveAlgorithm, args...; kwargs...) where {DT <: ForwardDiff.Dual}
     primal_sol = linearsolve_forwarddiff_solve!(
         cache::DualLinearCache, getfield(cache, :linear_cache).alg, args...; kwargs...)
-    dual_sol = linearsolve_dual_solution(getfield(cache,:linear_cache).u, getfield(cache, :rhs_list), cache)
+    dual_sol = linearsolve_dual_solution(primal_sol.u, getfield(cache, :rhs_list), cache)
 
     # For scalars, we still need to assign since cache.dual_u might not be pre-allocated
     if !(getfield(cache, :dual_u) isa AbstractArray)