Merge pull request #806 from ChrisRackauckas-Claude/fix-componentarray-forwarddiff-restructure

Fix ComponentArray type preservation in ForwardDiff specialization

Author: ChrisRackauckas

Project.toml

@@ -91,6 +91,7 @@ CUSOLVERRF = "0.2.6"
 ChainRulesCore = "1.25"
 CliqueTrees = "1.11.0"
 ConcreteStructs = "0.2.3"
+ComponentArrays = "0.15.29"
 DocStringExtensions = "0.9.3"
 EnumX = "1.0.4"
 EnzymeCore = "0.8.5"
@@ -148,6 +149,7 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 BlockDiagonals = "0a1fb500-61f7-11e9-3c65-f5ef3456f9f0"
 CliqueTrees = "60701a23-6482-424a-84db-faee86b9b1f8"
+ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
 FastAlmostBandedMatrices = "9d29842c-ecb8-4973-b1e9-a27b1157504e"
 FastLapackInterface = "29a986be-02c6-4525-aec4-84b980013641"
@@ -176,4 +178,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "Test", "IterativeSolvers", "InteractiveUtils", "KrylovKit", "KrylovPreconditioners", "Pkg", "Random", "SafeTestsets", "MultiFloats", "ForwardDiff", "HYPRE", "MPI", "BlockDiagonals", "FiniteDiff", "BandedMatrices", "FastAlmostBandedMatrices", "StaticArrays", "AllocCheck", "StableRNGs", "Zygote", "RecursiveFactorization", "Sparspak", "CliqueTrees", "FastLapackInterface", "SparseArrays", "ExplicitImports"]
+test = ["Aqua", "Test", "IterativeSolvers", "InteractiveUtils", "KrylovKit", "KrylovPreconditioners", "Pkg", "Random", "SafeTestsets", "MultiFloats", "ForwardDiff", "HYPRE", "MPI", "BlockDiagonals", "FiniteDiff", "BandedMatrices", "FastAlmostBandedMatrices", "StaticArrays", "AllocCheck", "StableRNGs", "Zygote", "RecursiveFactorization", "Sparspak", "CliqueTrees", "FastLapackInterface", "SparseArrays", "ExplicitImports", "ComponentArrays"]

ext/LinearSolveForwardDiffExt.jl

@@ -1,13 +1,15 @@
 module LinearSolveForwardDiffExt
 
 using LinearSolve
-using LinearSolve: SciMLLinearSolveAlgorithm, __init, LinearVerbosity, DefaultLinearSolver, DefaultAlgorithmChoice, defaultalg
+using LinearSolve: SciMLLinearSolveAlgorithm, __init, LinearVerbosity, DefaultLinearSolver,
+                   DefaultAlgorithmChoice, defaultalg
 using LinearAlgebra
 using ForwardDiff
 using ForwardDiff: Dual, Partials
 using SciMLBase
 using RecursiveArrayTools
 using SciMLLogging
+using ArrayInterface
 
 const DualLinearProblem = LinearProblem{
     <:Union{Number, <:AbstractArray, Nothing}, iip,
@@ -63,7 +65,7 @@ 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_u_cache .= cache.linear_cache.u
     cache.primal_b_cache .= cache.linear_cache.b
@@ -165,26 +167,27 @@ function linearsolve_dual_solution(
 end
 
 function linearsolve_dual_solution(u::AbstractArray, partials,
-        cache::DualLinearCache{DT}) where {T, V, N, DT <: Dual{T,V,N}}
+        cache::DualLinearCache{DT}) where {T, V, N, DT <: Dual{T, V, N}}
     # Optimized in-place version that reuses cache.dual_u
     linearsolve_dual_solution!(getfield(cache, :dual_u), u, partials)
     return getfield(cache, :dual_u)
 end
 
-function linearsolve_dual_solution!(dual_u::AbstractArray{DT}, u::AbstractArray, partials) where {T, V, N, DT <: Dual{T,V,N}}
+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
             V(partials[j][i])
         end
-        
+
         # Construct dual number in-place
-        dual_u[i] = DT(u[i], Partials{N,V}(partial_vals))
+        dual_u[i] = DT(u[i], Partials{N, V}(partial_vals))
     end
-    
+
     return dual_u
 end
 
@@ -279,7 +282,7 @@ function __dual_init(
         true,  # Cache is initially valid
         A,
         b,
-        zeros(dual_type, length(b))
+        ArrayInterface.restructure(b, zeros(dual_type, length(b)))
     )
 end
 
@@ -288,18 +291,20 @@ function SciMLBase.solve!(cache::DualLinearCache, args...; kwargs...)
 end
 
 function SciMLBase.solve!(
-        cache::DualLinearCache{DT}, alg::SciMLLinearSolveAlgorithm, args...; kwargs...) where {DT <: ForwardDiff.Dual}
+        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(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
     if !(getfield(cache, :dual_u) isa AbstractArray)
         setfield!(cache, :dual_u, dual_sol)
     end
 
     return SciMLBase.build_linear_solution(
-        getfield(cache, :linear_cache).alg, getfield(cache, :dual_u), primal_sol.resid, cache; primal_sol.retcode, primal_sol.iters, primal_sol.stats
+        getfield(cache, :linear_cache).alg, getfield(cache, :dual_u), primal_sol.resid, cache;
+        primal_sol.retcode, primal_sol.iters, primal_sol.stats
     )
 end
 

test/forwarddiff_overloads.jl

@@ -2,6 +2,7 @@ using LinearSolve
 using ForwardDiff
 using Test
 using SparseArrays
+using ComponentArrays
 
 function h(p)
     (A = [p[1] p[2]+1 p[2]^3;
@@ -194,4 +195,43 @@ A, b = h([ForwardDiff.Dual(5.0, 1.0, 0.0), ForwardDiff.Dual(5.0, 0.0, 1.0)])
 prob = LinearProblem(A, b)
 @test init(prob, GenericLUFactorization()) isa LinearSolve.LinearCache
 
-@test init(prob) isa LinearSolve.LinearCache
+@test init(prob) isa LinearSolve.LinearCache
+
+# Test ComponentArray with ForwardDiff (Issue SciML/DifferentialEquations.jl#1110)
+# This tests that ArrayInterface.restructure preserves ComponentArray structure
+
+# Direct test: ComponentVector with Dual elements should preserve structure
+ca_dual = ComponentArray(
+    a = ForwardDiff.Dual(1.0, 1.0, 0.0),
+    b = ForwardDiff.Dual(2.0, 0.0, 1.0)
+)
+A_dual = [ca_dual.a 1.0; 1.0 ca_dual.b]
+b_dual = ComponentArray(x = ca_dual.a + 1, y = ca_dual.b * 2)
+
+prob_dual = LinearProblem(A_dual, b_dual)
+sol_dual = solve(prob_dual)
+
+# The solution should preserve ComponentArray type
+@test sol_dual.u isa ComponentVector
+@test hasproperty(sol_dual.u, :x)
+@test hasproperty(sol_dual.u, :y)
+
+# Test gradient computation with ComponentArray inside ForwardDiff
+function component_linsolve(p)
+    # Create a matrix that depends on p
+    A = [p[1] p[2]; p[2] p[1] + 5]
+    # Create a ComponentArray RHS that depends on p
+    b_vec = ComponentArray(x = p[1] + 1, y = p[2] * 2)
+    prob = LinearProblem(A, b_vec)
+    sol = solve(prob)
+    # Return sum of solution
+    return sum(sol.u)
+end
+
+p_test = [2.0, 3.0]
+# This will internally create Dual numbers and ComponentArrays with Dual elements
+grad = ForwardDiff.gradient(component_linsolve, p_test)
+@test grad isa Vector
+@test length(grad) == 2
+@test !any(isnan, grad)
+@test !any(isinf, grad)