Skip to content
This repository was archived by the owner on Aug 22, 2025. It is now read-only.

Commit db2f2ae

Browse files
ChrisRackauckasChrisRackauckas
authored
Merge pull request #253 from avik-pal/ap/autosparse
New High Level Interface for Sparse Jacobian Computation
2 parents 9c458cd + 713e6ee commit db2f2ae

32 files changed

+1169
-486
lines changed

‎.gitignore‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,3 +6,5 @@ Manifest.toml
66

77
docs/build/
88
docs/site/
9+
10+
.vscode

‎Project.toml‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
name = "SparseDiffTools"
22
uuid = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
33
authors = ["Pankaj Mishra <pankajmishra1511@gmail.com>", "Chris Rackauckas <contact@chrisrackauckas.com>"]
4-
version = "2.4.1"
4+
version = "2.5.0"
55

66
[deps]
77
ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -13,38 +13,45 @@ FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
1313
ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
1414
Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
1515
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
16+
PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930"
1617
Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
17-
Requires = "ae029012-a4dd-5104-9daa-d747884805df"
1818
SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
1919
Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
2020
SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
2121
StaticArrayInterface = "0d7ed370-da01-4f52-bd93-41d350b8b718"
2222
StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
2323
Tricks = "410a4b4d-49e4-4fbc-ab6d-cb71b17b3775"
24+
UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
2425
VertexSafeGraphs = "19fa3120-7c27-5ec5-8db8-b0b0aa330d6f"
2526

2627
[weakdeps]
28+
Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
29+
Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
2730
Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
2831

2932
[extensions]
33+
SparseDiffToolsEnzymeExt = "Enzyme"
34+
SparseDiffToolsSymbolicsExt = "Symbolics"
3035
SparseDiffToolsZygoteExt = "Zygote"
3136

3237
[compat]
33-
ADTypes = "0.1"
38+
ADTypes = "0.2.1"
3439
Adapt = "3.0"
3540
ArrayInterface = "7.4.2"
3641
Compat = "4"
3742
DataStructures = "0.18"
43+
Enzyme = "0.11"
3844
FiniteDiff = "2.8.1"
3945
ForwardDiff = "0.10"
4046
Graphs = "1"
4147
Reexport = "1"
42-
Requires = "1"
4348
SciMLOperators = "0.2.11, 0.3"
4449
Setfield = "1"
4550
StaticArrayInterface = "1.3"
4651
StaticArrays = "1"
52+
Symbolics = "5.5"
4753
Tricks = "0.1.6"
54+
UnPack = "1"
4855
VertexSafeGraphs = "0.2"
4956
Zygote = "0.6"
5057
julia = "1.6"
@@ -54,6 +61,7 @@ ArrayInterfaceBandedMatrices = "2e50d22c-5be1-4042-81b1-c572ed69783d"
5461
ArrayInterfaceBlockBandedMatrices = "5331f1e9-51c7-46b0-a9b0-df4434785e0a"
5562
BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
5663
BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
64+
Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
5765
IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
5866
Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
5967
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -64,4 +72,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
6472
Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
6573

6674
[targets]
67-
test = ["Test", "ArrayInterfaceBandedMatrices", "ArrayInterfaceBlockBandedMatrices", "BandedMatrices", "BlockBandedMatrices", "IterativeSolvers", "Pkg", "Random", "SafeTestsets", "Symbolics", "Zygote", "StaticArrays"]
75+
test = ["Test", "ArrayInterfaceBandedMatrices", "ArrayInterfaceBlockBandedMatrices", "BandedMatrices", "BlockBandedMatrices", "Enzyme", "IterativeSolvers", "Pkg", "Random", "SafeTestsets", "Symbolics", "Zygote", "StaticArrays"]

‎README.md‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,67 @@ function g(x) # out-of-place
4747
end
4848
```
4949

50+
## High Level API
51+
52+
We need to perform the following steps to utilize SparseDiffTools:
53+
54+
1. Specify a Sparsity Detection Algorithm. There are 3 possible choices currently:
55+
1. `NoSparsityDetection`: This will ignore any AD choice and compute the dense Jacobian
56+
2. `JacPrototypeSparsityDetection`: If you already know the sparsity pattern, you can
57+
specify it as `JacPrototypeSparsityDetection(; jac_prototype=<sparsity pattern>)`.
58+
3. `SymbolicsSparsityDetection`: This will use `Symbolics.jl` to automatically detect
59+
the sparsity pattern. (Note that `Symbolics.jl` must be explicitly loaded before
60+
using this functionality.)
61+
2. Now choose an AD backend from `ADTypes.jl`:
62+
1. If using a Non `*Sparse*` type, then we will not use sparsity detection.
63+
2. All other sparse AD types will internally compute the proper sparsity pattern, and
64+
try to exploit that.
65+
3. Now there are 2 options:
66+
1. Precompute the cache using `sparse_jacobian_cache` and use the `sparse_jacobian` or
67+
`sparse_jacobian!` functions to compute the Jacobian. This option is recommended if
68+
you are repeatedly computing the Jacobian for the same function.
69+
2. Directly use `sparse_jacobian` or `sparse_jacobian!` to compute the Jacobian. This
70+
option should be used if you are only computing the Jacobian once.
71+
72+
```julia
73+
using Symbolics
74+
75+
sd = SymbolicsSparsityDetection()
76+
adtype = AutoSparseFiniteDiff()
77+
x = rand(30)
78+
y = similar(x)
79+
80+
# Option 1
81+
## OOP Function
82+
cache = sparse_jacobian_cache(adtype, sd, g, x; fx=y) # Passing `fx` is needed if size(y) != size(x)
83+
J = sparse_jacobian(adtype, cache, g, x)
84+
### Or
85+
J_preallocated = similar(J)
86+
sparse_jacobian!(J_preallocated, adtype, cache, g, x)
87+
88+
## IIP Function
89+
cache = sparse_jacobian_cache(adtype, sd, f, y, x)
90+
J = sparse_jacobian(adtype, cache, f, y, x)
91+
### Or
92+
J_preallocated = similar(J)
93+
sparse_jacobian!(J_preallocated, adtype, cache, f, y, x)
94+
95+
# Option 2
96+
## OOP Function
97+
J = sparse_jacobian(adtype, sd, g, x)
98+
### Or
99+
J_preallocated = similar(J)
100+
sparse_jacobian!(J_preallocated, adtype, sd, g, x)
101+
102+
## IIP Function
103+
J = sparse_jacobian(adtype, sd, f, y, x)
104+
### Or
105+
J_preallocated = similar(J)
106+
sparse_jacobian!(J_preallocated, adtype, sd, f, y, x)
107+
```
108+
109+
## Lower Level API
110+
50111
For this function, we know that the sparsity pattern of the Jacobian is a
51112
`Tridiagonal` matrix. However, if we didn't know the sparsity pattern for
52113
the Jacobian, we could use the `Symbolics.jacobian_sparsity` function to automatically

‎docs/make.jl‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,13 @@ using Documenter, SparseDiffTools
33
include("pages.jl")
44

55
makedocs(sitename = "SparseDiffTools.jl",
6-
authors = "Chris Rackauckas",
7-
modules = [SparseDiffTools],
8-
clean = true,
9-
doctest = false,
10-
format = Documenter.HTML(assets = ["assets/favicon.ico"],
11-
canonical = "https://docs.sciml.ai/SparseDiffTools/stable/"),
12-
pages = pages)
6+
authors = "Chris Rackauckas",
7+
modules = [SparseDiffTools],
8+
clean = true,
9+
doctest = false,
10+
format = Documenter.HTML(assets = ["assets/favicon.ico"],
11+
canonical = "https://docs.sciml.ai/SparseDiffTools/stable/"),
12+
pages = pages)
1313

1414
deploydocs(repo = "github.com/JuliaDiff/SparseDiffTools.jl.git";
15-
push_preview = true)
15+
push_preview = true)

‎docs/src/index.md‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,67 @@ function g(x) # out-of-place
3737
end
3838
```
3939

40+
## High Level API
41+
42+
We need to perform the following steps to utilize SparseDiffTools:
43+
44+
1. Specify a Sparsity Detection Algorithm. There are 3 possible choices currently:
45+
1. `NoSparsityDetection`: This will ignore any AD choice and compute the dense Jacobian
46+
2. `JacPrototypeSparsityDetection`: If you already know the sparsity pattern, you can
47+
specify it as `JacPrototypeSparsityDetection(; jac_prototype=<sparsity pattern>)`.
48+
3. `SymbolicsSparsityDetection`: This will use `Symbolics.jl` to automatically detect
49+
the sparsity pattern. (Note that `Symbolics.jl` must be explicitly loaded before
50+
using this functionality.)
51+
2. Now choose an AD backend from `ADTypes.jl`:
52+
1. If using a Non `*Sparse*` type, then we will not use sparsity detection.
53+
2. All other sparse AD types will internally compute the proper sparsity pattern, and
54+
try to exploit that.
55+
3. Now there are 2 options:
56+
1. Precompute the cache using `sparse_jacobian_cache` and use the `sparse_jacobian` or
57+
`sparse_jacobian!` functions to compute the Jacobian. This option is recommended if
58+
you are repeatedly computing the Jacobian for the same function.
59+
2. Directly use `sparse_jacobian` or `sparse_jacobian!` to compute the Jacobian. This
60+
option should be used if you are only computing the Jacobian once.
61+
62+
```julia
63+
using Symbolics
64+
65+
sd = SymbolicsSparsityDetection()
66+
adtype = AutoSparseFiniteDiff()
67+
x = rand(30)
68+
y = similar(x)
69+
70+
# Option 1
71+
## OOP Function
72+
cache = sparse_jacobian_cache(adtype, sd, g, x; fx=y) # Passing `fx` is needed if size(y) != size(x)
73+
J = sparse_jacobian(adtype, cache, g, x)
74+
### Or
75+
J_preallocated = similar(J)
76+
sparse_jacobian!(J_preallocated, adtype, cache, g, x)
77+
78+
## IIP Function
79+
cache = sparse_jacobian_cache(adtype, sd, f, y, x)
80+
J = sparse_jacobian(adtype, cache, f, y, x)
81+
### Or
82+
J_preallocated = similar(J)
83+
sparse_jacobian!(J_preallocated, adtype, cache, f, y, x)
84+
85+
# Option 2
86+
## OOP Function
87+
J = sparse_jacobian(adtype, sd, g, x)
88+
### Or
89+
J_preallocated = similar(J)
90+
sparse_jacobian!(J_preallocated, adtype, sd, g, x)
91+
92+
## IIP Function
93+
J = sparse_jacobian(adtype, sd, f, y, x)
94+
### Or
95+
J_preallocated = similar(J)
96+
sparse_jacobian!(J_preallocated, adtype, sd, f, y, x)
97+
```
98+
99+
## Lower Level API
100+
40101
For this function, we know that the sparsity pattern of the Jacobian is a
41102
`Tridiagonal` matrix. However, if we didn't know the sparsity pattern for
42103
the Jacobian, we could use the `Symbolics.jacobian_sparsity` function to automatically

‎ext/SparseDiffToolsEnzymeExt.jl‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
module SparseDiffToolsEnzymeExt
2+
3+
import ArrayInterface: fast_scalar_indexing
4+
import SparseDiffTools: __f̂,
5+
__maybe_copy_x, __jacobian!, __gradient, __gradient!, AutoSparseEnzyme
6+
# FIXME: For Enzyme we currently assume reverse mode
7+
import ADTypes: AutoEnzyme
8+
using Enzyme
9+
10+
using ForwardDiff
11+
12+
## Satisfying High-Level Interface for Sparse Jacobians
13+
function __gradient(::Union{AutoSparseEnzyme, AutoEnzyme}, f, x, cols)
14+
dx = zero(x)
15+
autodiff(Reverse, __f̂, Const(f), Duplicated(x, dx), Const(cols))
16+
return vec(dx)
17+
end
18+
19+
function __gradient!(::Union{AutoSparseEnzyme, AutoEnzyme}, f!, fx, x, cols)
20+
dx = zero(x)
21+
dfx = zero(fx)
22+
autodiff(Reverse, __f̂, Active, Const(f!), Duplicated(fx, dfx), Duplicated(x, dx),
23+
Const(cols))
24+
return dx
25+
end
26+
27+
function __jacobian!(J::AbstractMatrix, ::Union{AutoSparseEnzyme, AutoEnzyme}, f, x)
28+
J .= jacobian(Reverse, f, x, Val(size(J, 1)))
29+
return J
30+
end
31+
32+
@views function __jacobian!(J, ad::Union{AutoSparseEnzyme, AutoEnzyme}, f!, fx, x)
33+
# This version is slowish not sure how to do jacobians for inplace functions
34+
@warn "Current code for computing jacobian for inplace functions in Enzyme is slow." maxlog=1
35+
dfx = zero(fx)
36+
dx = zero(x)
37+
38+
function __f_row_idx(f!, fx, x, row_idx)
39+
f!(fx, x)
40+
if fast_scalar_indexing(fx)
41+
return fx[row_idx]
42+
else
43+
# Avoid the GPU Arrays scalar indexing error
44+
return sum(selectdim(fx, 1, row_idx:row_idx))
45+
end
46+
end
47+
48+
for row_idx in 1:size(J, 1)
49+
autodiff(Reverse, __f_row_idx, Const(f!), DuplicatedNoNeed(fx, dfx),
50+
Duplicated(x, dx), Const(row_idx))
51+
J[row_idx, :] .= dx
52+
fill!(dfx, 0)
53+
fill!(dx, 0)
54+
end
55+
56+
return J
57+
end
58+
59+
__maybe_copy_x(::Union{AutoSparseEnzyme, AutoEnzyme}, x::SubArray) = copy(x)
60+
61+
end

‎ext/SparseDiffToolsSymbolicsExt.jl‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
module SparseDiffToolsSymbolicsExt
2+
3+
using SparseDiffTools, Symbolics
4+
import SparseDiffTools: AbstractSparseADType
5+
6+
function (alg::SymbolicsSparsityDetection)(ad::AbstractSparseADType, f, x; fx = nothing,
7+
kwargs...)
8+
fx = fx === nothing ? similar(f(x)) : dx
9+
f!(y, x) = (y .= f(x))
10+
J = Symbolics.jacobian_sparsity(f!, fx, x)
11+
_alg = JacPrototypeSparsityDetection(J, alg.alg)
12+
return _alg(ad, f, x; fx, kwargs...)
13+
end
14+
15+
function (alg::SymbolicsSparsityDetection)(ad::AbstractSparseADType, f!, fx, x; kwargs...)
16+
J = Symbolics.jacobian_sparsity(f!, fx, x)
17+
_alg = JacPrototypeSparsityDetection(J, alg.alg)
18+
return _alg(ad, f!, fx, x; kwargs...)
19+
end
20+
21+
end

0 commit comments

Comments
 (0)