Skip to content

Commit efa1c67

Browse files
avik-palavik-pal
committed
Documentation for AD Backends
1 parent c9f7d3c commit efa1c67

File tree

5 files changed

+206
-1
lines changed

5 files changed

+206
-1
lines changed

‎docs/Project.toml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
[deps]
2+
ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
23
AlgebraicMultigrid = "2169fc97-5a83-5252-b627-83903c6c433c"
34
ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
45
BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
@@ -22,6 +23,7 @@ Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4"
2223
Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
2324

2425
[compat]
26+
ADTypes = "0.2.6"
2527
AlgebraicMultigrid = "0.5, 0.6"
2628
ArrayInterface = "6, 7"
2729
BenchmarkTools = "1"

‎docs/pages.jl

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ pages = [
1313
"basics/nonlinear_functions.md",
1414
"basics/solve.md",
1515
"basics/nonlinear_solution.md",
16+
"basics/autodiff.md",
1617
"basics/termination_condition.md",
1718
"basics/diagnostics_api.md",
1819
"basics/sparsity_detection.md",

‎docs/src/basics/autodiff.md

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
# Automatic Differentiation Backends
2+
3+
## Summary of Finite Differencing Backends
4+
5+
- [`AutoFiniteDiff`](@ref): Finite differencing, not optimal but always applicable.
6+
- [`AutoSparseFiniteDiff`](@ref): Sparse version of [`AutoFiniteDiff`](@ref).
7+
8+
## Summary of Forward Mode AD Backends
9+
10+
- [`AutoForwardDiff`](@ref): The best choice for dense problems.
11+
- [`AutoSparseForwardDiff`](@ref): Sparse version of [`AutoForwardDiff`](@ref).
12+
- [`AutoPolyesterForwardDiff`](@ref): Might be faster than [`AutoForwardDiff`](@ref) for
13+
large problems. Requires `PolyesterForwardDiff.jl` to be installed and loaded.
14+
15+
## Summary of Reverse Mode AD Backends
16+
17+
- [`AutoZygote`](@ref): The fastest choice for non-mutating array-based (BLAS) functions.
18+
- [`AutoSparseZygote`](@ref): Sparse version of [`AutoZygote`](@ref).
19+
- [`AutoEnzyme`](@ref): Uses `Enzyme.jl` Reverse Mode and should be considered
20+
experimental.
21+
22+
!!! note
23+
24+
If `PolyesterForwardDiff.jl` is installed and loaded, then `SimpleNonlinearSolve.jl`
25+
will automatically use `AutoPolyesterForwardDiff` as the default AD backend.
26+
27+
## API Reference
28+
29+
### Finite Differencing Backends
30+
31+
```@docs
32+
AutoFiniteDiff
33+
AutoSparseFiniteDiff
34+
```
35+
36+
### Forward Mode AD Backends
37+
38+
```@docs
39+
AutoForwardDiff
40+
AutoSparseForwardDiff
41+
AutoPolyesterForwardDiff
42+
```
43+
44+
### Reverse Mode AD Backends
45+
46+
```@docs
47+
AutoZygote
48+
AutoSparseZygote
49+
AutoEnzyme
50+
NonlinearSolve.AutoSparseEnzyme
51+
```

‎src/NonlinearSolve.jl

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ import PrecompileTools: @recompile_invalidations, @compile_workload, @setup_work
2424

2525
import SciMLBase: AbstractNonlinearAlgorithm, JacobianWrapper, AbstractNonlinearProblem,
2626
AbstractSciMLOperator, NLStats, _unwrap_val, has_jac, isinplace
27-
import SparseDiffTools: AbstractSparsityDetection
27+
import SparseDiffTools: AbstractSparsityDetection, AutoSparseEnzyme
2828
import StaticArraysCore: StaticArray, SVector, SArray, MArray, Size, SMatrix, MMatrix
2929
end
3030

@@ -40,6 +40,7 @@ const True = Val(true)
4040
const False = Val(false)
4141

4242
include("abstract_types.jl")
43+
include("adtypes.jl")
4344
include("timer_outputs.jl")
4445
include("internal/helpers.jl")
4546

‎src/adtypes.jl

Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,150 @@
1+
# This just documents the AD types from ADTypes.jl
2+
3+
"""
4+
AutoFiniteDiff(; fdtype = Val(:forward), fdjtype = fdtype, fdhtype = Val(:hcentral))
5+
6+
This uses [FiniteDiff.jl](https://github.com/JuliaDiff/FiniteDiff.jl). While not necessarily
7+
the most efficient, this is the only choice that doesn't require the `f` function to be
8+
automatically differentiable, which means it applies to any choice. However, because it's
9+
using finite differencing, one needs to be careful as this procedure introduces numerical
10+
error into the derivative estimates.
11+
12+
- Compatible with GPUs
13+
- Can be used for Jacobian-Vector Products (JVPs)
14+
- Can be used for Vector-Jacobian Products (VJPs)
15+
- Supports both inplace and out-of-place functions
16+
17+
### Keyword Arguments
18+
19+
- `fdtype`: the method used for defining the gradient
20+
- `fdjtype`: the method used for defining the Jacobian of constraints.
21+
- `fdhtype`: the method used for defining the Hessian
22+
"""
23+
AutoFiniteDiff
24+
25+
"""
26+
AutoSparseFiniteDiff()
27+
28+
Sparse Version of [`AutoFiniteDiff`](@ref) that uses
29+
[FiniteDiff.jl](https://github.com/JuliaDiff/FiniteDiff.jl) and the column color vector of
30+
the Jacobian Matrix to efficiently compute the Sparse Jacobian.
31+
32+
- Supports both inplace and out-of-place functions
33+
"""
34+
AutoSparseFiniteDiff
35+
36+
"""
37+
AutoForwardDiff(; chunksize = nothing, tag = nothing)
38+
AutoForwardDiff{chunksize, tagType}(tag::tagType)
39+
40+
This uses the [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl) package. It is
41+
the fastest choice for square or wide systems. It is easy to use and compatible with most
42+
Julia functions which have loose type restrictions.
43+
44+
- Compatible with GPUs
45+
- Can be used for Jacobian-Vector Products (JVPs)
46+
- Supports both inplace and out-of-place functions
47+
48+
For type-stability of internal operations, a positive `chunksize` must be provided.
49+
50+
### Keyword Arguments
51+
52+
- `chunksize`: Count of dual numbers that can be propagated simultaneously. Setting this
53+
number to a high value will lead to slowdowns. Use
54+
[`NonlinearSolve.pickchunksize`](@ref) to get a proper value.
55+
- `tag`: Used to avoid perturbation confusion. If set to `nothing`, we use a custom tag.
56+
"""
57+
AutoForwardDiff
58+
59+
"""
60+
AutoSparseForwardDiff(; chunksize = nothing, tag = nothing)
61+
AutoSparseForwardDiff{chunksize, tagType}(tag::tagType)
62+
63+
Sparse Version of [`AutoForwardDiff`](@ref) that uses
64+
[ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl) and the column color vector of
65+
the Jacobian Matrix to efficiently compute the Sparse Jacobian.
66+
67+
- Supports both inplace and out-of-place functions
68+
69+
For type-stability of internal operations, a positive `chunksize` must be provided.
70+
71+
### Keyword Arguments
72+
73+
- `chunksize`: Count of dual numbers that can be propagated simultaneously. Setting this
74+
number to a high value will lead to slowdowns. Use
75+
[`NonlinearSolve.pickchunksize`](@ref) to get a proper value.
76+
- `tag`: Used to avoid perturbation confusion. If set to `nothing`, we use a custom tag.
77+
"""
78+
AutoSparseForwardDiff
79+
80+
@static if isdefined(ADTypes, :AutoPolyesterForwardDiff)
81+
"""
82+
AutoPolyesterForwardDiff(; chunksize = nothing)
83+
84+
Uses [`PolyesterForwardDiff.jl`](https://github.com/JuliaDiff/PolyesterForwardDiff.jl)
85+
to compute the jacobian. This is essentially parallelized `ForwardDiff.jl`.
86+
87+
- Supports both inplace and out-of-place functions
88+
89+
### Keyword Arguments
90+
91+
- `chunksize`: Count of dual numbers that can be propagated simultaneously. Setting
92+
this number to a high value will lead to slowdowns. Use
93+
[`NonlinearSolve.pickchunksize`](@ref) to get a proper value.
94+
"""
95+
AutoPolyesterForwardDiff
96+
end
97+
98+
"""
99+
AutoZygote()
100+
101+
Uses [`Zygote.jl`](https://github.com/FluxML/Zygote.jl) package. This is the staple
102+
reverse-mode AD that handles a large portion of Julia with good efficiency.
103+
104+
- Compatible with GPUs
105+
- Can be used for Vector-Jacobian Products (VJPs)
106+
- Supports only out-of-place functions
107+
108+
For VJPs this is the current best choice. This is the most efficient method for long
109+
jacobians.
110+
"""
111+
AutoZygote
112+
113+
"""
114+
AutoSparseZygote()
115+
116+
Sparse version of [`AutoZygote`](@ref) that uses
117+
[`Zygote.jl`](https://github.com/FluxML/Zygote.jl) and the row color vector of
118+
the Jacobian Matrix to efficiently compute the Sparse Jacobian.
119+
120+
- Supports only out-of-place functions
121+
122+
This is efficient only for long jacobians or if the maximum value of the row color vector is
123+
significantly lower than the maximum value of the column color vector.
124+
"""
125+
AutoSparseZygote
126+
127+
"""
128+
AutoEnzyme()
129+
130+
Uses reverse mode [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl). This is currently
131+
experimental, and not extensively tested on our end. We only support Jacobian construction
132+
and VJP support is currently not implemented.
133+
134+
- Supports both inplace and out-of-place functions
135+
"""
136+
AutoEnzyme
137+
138+
"""
139+
AutoSparseEnzyme()
140+
141+
Sparse version of [`AutoEnzyme`](@ref) that uses
142+
[Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl) and the row color vector of
143+
the Jacobian Matrix to efficiently compute the Sparse Jacobian.
144+
145+
- Supports both inplace and out-of-place functions
146+
147+
This is efficient only for long jacobians or if the maximum value of the row color vector is
148+
significantly lower than the maximum value of the column color vector.
149+
"""
150+
AutoSparseEnzyme

0 commit comments

Comments
 (0)