From a34ba9bceeddbfbe61584d1344b655ddbd992739 Mon Sep 17 00:00:00 2001 From: jClugstor Date: Thu, 4 Dec 2025 21:08:36 -0500 Subject: [PATCH 1/2] initialize parts of cache based on size of u instead of b --- ext/LinearSolveForwardDiffExt.jl | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/ext/LinearSolveForwardDiffExt.jl b/ext/LinearSolveForwardDiffExt.jl index 31bfab6b2..c6dfb2cc6 100644 --- a/ext/LinearSolveForwardDiffExt.jl +++ b/ext/LinearSolveForwardDiffExt.jl @@ -76,9 +76,8 @@ function linearsolve_forwarddiff_solve!(cache::DualLinearCache, alg, args...; kw ∂_b = cache.partials_b xp_linsolve_rhs!(uu, ∂_A, ∂_b, cache) - + rhs_list = cache.rhs_list - cache.linear_cache.u .= cache.dual_u0_cache # We can reuse the linear cache, because the same factorization will work for the partials. for i in eachindex(rhs_list) @@ -177,7 +176,6 @@ function linearsolve_dual_solution!(dual_u::AbstractArray{DT}, u::AbstractArray, partials) where {T, V, N, DT <: Dual{T, V, N}} # Direct in-place construction of dual numbers without temporary allocations n_partials = length(partials) - for i in eachindex(u, dual_u) # Extract partials for this element directly partial_vals = ntuple(Val(N)) do j @@ -281,13 +279,13 @@ function __dual_init( partials_A_list, partials_b_list, rhs_list, - similar(new_b), - similar(new_b), - similar(new_b), + similar(non_partial_cache.u), # Use u's size, not b's size + similar(non_partial_cache.u), # primal_u_cache + similar(new_b), # primal_b_cache true, # Cache is initially valid A, b, - ArrayInterface.restructure(b, zeros(dual_type, length(b))) + ArrayInterface.restructure(non_partial_cache.u, zeros(dual_type, length(non_partial_cache.u))) ) end @@ -300,6 +298,7 @@ function SciMLBase.solve!( 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) # For scalars, we still need to assign since cache.dual_u might not be pre-allocated From 9fe1907979500fd05f1472f623153df6e27c3d16 Mon Sep 17 00:00:00 2001 From: jClugstor Date: Fri, 5 Dec 2025 10:39:54 -0500 Subject: [PATCH 2/2] account for restructuring based on b, or u based on if b matches u --- ext/LinearSolveForwardDiffExt.jl | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/ext/LinearSolveForwardDiffExt.jl b/ext/LinearSolveForwardDiffExt.jl index c6dfb2cc6..018a9b8cd 100644 --- a/ext/LinearSolveForwardDiffExt.jl +++ b/ext/LinearSolveForwardDiffExt.jl @@ -270,6 +270,13 @@ function __dual_init( else rhs_list = nothing end + # Use b for restructuring if sizes match (square system), otherwise use u (non-square) + # This preserves ComponentArray structure from b when possible + dual_u_init = if length(non_partial_cache.u) == length(b) + ArrayInterface.restructure(b, zeros(dual_type, length(b))) + else + ArrayInterface.restructure(non_partial_cache.u, zeros(dual_type, length(non_partial_cache.u))) + end return DualLinearCache{dual_type}( non_partial_cache, @@ -285,7 +292,7 @@ function __dual_init( true, # Cache is initially valid A, b, - ArrayInterface.restructure(non_partial_cache.u, zeros(dual_type, length(non_partial_cache.u))) + dual_u_init ) end