Skip to content

Specialize split dual AD path on sparse matrix in ForwardDiff extension #1048

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ChrisRackauckas merged 5 commits into from
Jun 18, 2026
+31 −4
Merged

Specialize split dual AD path on sparse matrix in ForwardDiff extension #1048

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
ee8884f
Make partials list dispatch signatures consistent
hersle Jun 17, 2026
1dc68ec
Move SparseArrays to deps so it is available in the ForwardDiff exten…
hersle Jun 16, 2026
0abfd72
Specialize update_partials_list! and partials_to_list for SparseMatri…
hersle Jun 16, 2026
a999104
Handle changing sparsity pattern in update_partials_list! for SparseM…
hersle Jun 17, 2026
0492ab1
Handle SparseMatrixCSC with empty nonzeros in partials_to_list
hersle Jun 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
SciMLLogging = "a6db7da4-7206-11f0-1eab-35f2a5dbe1d1"
SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
SparseColumnPivotedQR = "a57abbd0-fea5-4d57-96be-5e525945e8e4"
StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"

Expand Down Expand Up @@ -62,7 +63,6 @@ PartitionedSolvers = "11b65f7f-80ac-401b-9ef2-3db765482d62"
PureUMFPACK = "b7e1f0a2-3c4d-4e5f-9a0b-1c2d3e4f5a6b"
RecursiveFactorization = "f2c3362d-daeb-58d1-803e-2bc74f2840b4"
STRUMPACK_jll = "86fbd0b9-476f-557c-b766-62c724b42d8c"
SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
SparseMatricesCSR = "a0a7dd2c-ebf4-11e9-1f05-cf50bc540ca1"
Sparspak = "e56a9233-b9d6-4f03-8d0f-1825330902ac"
SpecializingFactorizations = "fa08b7a1-13d3-4faf-875d-5cbc1520e3f3"
Expand Down
33 changes: 30 additions & 3 deletions ext/LinearSolveForwardDiffExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ using LinearSolve
using LinearSolve: SciMLLinearSolveAlgorithm, __init, LinearVerbosity, DefaultLinearSolver,
DefaultAlgorithmChoice, defaultalg, reinit!
using LinearAlgebra
using SparseArrays
using ForwardDiff
using ForwardDiff: Dual, Partials
using SciMLBase
Expand Down Expand Up @@ -636,10 +637,9 @@ function update_partials_list!(partial_matrix::AbstractVector{T}, list_cache) wh
return list_cache
end

function update_partials_list!(partial_matrix, list_cache)
function update_partials_list!(partial_matrix::AbstractMatrix{T}, list_cache) where {T}
p = length(first(partial_matrix))
m, n = size(partial_matrix)

for k in 1:p
for i in 1:m
for j in 1:n
Expand All @@ -655,7 +655,7 @@ function partials_to_list(partial_matrix::AbstractVector{T}) where {T}
return [[partial[i] for partial in partial_matrix] for i in p]
end

function partials_to_list(partial_matrix)
function partials_to_list(partial_matrix::AbstractMatrix{T}) where {T}
p = length(first(partial_matrix))
m, n = size(partial_matrix)
res_list = fill(zeros(typeof(partial_matrix[1, 1][1]), m, n), p)
Expand All @@ -671,4 +671,31 @@ function partials_to_list(partial_matrix)
return res_list
end

# Specializations for sparse matrices

function partials_to_list(partial_matrix::SparseMatrixCSC)
nz = nonzeros(partial_matrix)
m, n = size(partial_matrix)
T = eltype(partial_matrix)
p = ForwardDiff.npartials(T)
V = ForwardDiff.valtype(T) # use type for concrete array below in empty-nz case (e.g. all-zero Jacobian at init)
return [SparseMatrixCSC(m, n, copy(partial_matrix.colptr), copy(partial_matrix.rowval),
V[nz[i][k] for i in eachindex(nz)]) for k in 1:p]
end

function update_partials_list!(partial_matrix::SparseMatrixCSC, list_cache)
nz = nonzeros(partial_matrix)
if length(nz) != length(nonzeros(first(list_cache))) # TODO: more precise?
list_cache .= partials_to_list(partial_matrix) # sparsity pattern changed
else
for k in eachindex(list_cache)
nz_k = nonzeros(list_cache[k])
@inbounds for i in eachindex(nz, nz_k)
nz_k[i] = nz[i][k]
end
end
end
return list_cache
end

end
Loading