Make length of Partials known at compile time in ForwardDiff overloads (#727)

* make sure that Partials length is known

* add Dual problem JET tests

* Update test/nopre/jet.jl

* use dual_prob

---------

Co-authored-by: Christopher Rackauckas <[email protected]>

Author: jClugstor

ext/LinearSolveForwardDiffExt.jl

@@ -34,7 +34,7 @@ const DualBLinearProblem = LinearProblem{
 const DualAbstractLinearProblem = Union{
     DualLinearProblem, DualALinearProblem, DualBLinearProblem}
 
-LinearSolve.@concrete mutable struct DualLinearCache{DT <: Dual}
+LinearSolve.@concrete mutable struct DualLinearCache{DT}
     linear_cache
 
     partials_A
@@ -113,10 +113,10 @@ function linearsolve_dual_solution(
 end
 
 function linearsolve_dual_solution(u::AbstractArray, partials,
-        cache::DualLinearCache{DT}) where {DT}
+        cache::DualLinearCache{DT}) where {T, V, N, DT <: Dual{T,V,N}}
     # Handle single-level duals for arrays
     partials_list = RecursiveArrayTools.VectorOfArray(partials)
-    return map(((uᵢ, pᵢ),) -> DT(uᵢ, Partials(Tuple(pᵢ))),
+    return map(((uᵢ, pᵢ),) -> DT(uᵢ, Partials{N,V}(NTuple{N,V}(pᵢ))),
         zip(u, partials_list[i, :] for i in 1:length(partials_list.u[1])))
 end
 

test/nopre/jet.jl

@@ -1,4 +1,4 @@
-using LinearSolve, RecursiveFactorization, LinearAlgebra, SparseArrays, Test
+using LinearSolve, ForwardDiff, RecursiveFactorization, LinearAlgebra, SparseArrays, Test
 using JET
 
 # Dense problem setup
@@ -22,6 +22,18 @@ prob_sparse = LinearProblem(A_sparse, b)
 A_sparse_spd = sparse(A_spd)
 prob_sparse_spd = LinearProblem(A_sparse_spd, b)
 
+# Dual problem set up 
+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)])
+
+dual_prob = LinearProblem(A, b)
+
 @testset "JET Tests for Dense Factorizations" begin
     # Working tests - these pass JET optimization checks
     JET.@test_opt init(prob, nothing)
@@ -109,3 +121,11 @@ end
     JET.@test_opt solve(prob) broken=true
     JET.@test_opt solve(prob_sparse) broken=true
 end
+
+@testset "JET Tests for creating Dual solutions" begin
+    # Make sure there's no runtime dispatch when making solutions of Dual problems
+    dual_cache = init(dual_prob)
+    ext = Base.get_extension(LinearSolve, :LinearSolveForwardDiffExt)
+    JET.@test_opt ext.linearsolve_dual_solution(
+        [1.0, 1.0, 1.0], [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]], dual_cache)
+end