Merge pull request #354 from avik-pal/ap/polyester-forward-diff

AD Section in Documentation

.github/workflows/Documentation.yml

@@ -23,7 +23,7 @@ jobs:
           JULIA_DEBUG: "Documenter"
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} # For authentication with GitHub Actions token
           DOCUMENTER_KEY: ${{ secrets.DOCUMENTER_KEY }} # For authentication with SSH deploy key
-        run: julia --project=docs/ --code-coverage=user docs/make.jl
+        run: julia --project=docs/ --code-coverage=user --color=yes docs/make.jl
       - uses: julia-actions/julia-processcoverage@v1
       - uses: codecov/codecov-action@v4
         with:

docs/pages.jl

@@ -13,6 +13,7 @@ pages = [
         "basics/nonlinear_functions.md",
         "basics/solve.md",
         "basics/nonlinear_solution.md",
+        "basics/autodiff.md",
         "basics/termination_condition.md",
         "basics/diagnostics_api.md",
         "basics/sparsity_detection.md",

docs/src/basics/autodiff.md

@@ -0,0 +1,62 @@
+# Automatic Differentiation Backends
+
+## Summary of Finite Differencing Backends
+
+  - [`AutoFiniteDiff`](@ref): Finite differencing, not optimal but always applicable.
+  - [`AutoSparseFiniteDiff`](@ref): Sparse version of [`AutoFiniteDiff`](@ref).
+
+## Summary of Forward Mode AD Backends
+
+  - [`AutoForwardDiff`](@ref): The best choice for dense problems.
+  - [`AutoSparseForwardDiff`](@ref): Sparse version of [`AutoForwardDiff`](@ref).
+  - [`AutoPolyesterForwardDiff`](@ref): Might be faster than [`AutoForwardDiff`](@ref) for
+    large problems. Requires `PolyesterForwardDiff.jl` to be installed and loaded.
+
+## Summary of Reverse Mode AD Backends
+
+  - [`AutoZygote`](@ref): The fastest choice for non-mutating array-based (BLAS) functions.
+  - [`AutoSparseZygote`](@ref): Sparse version of [`AutoZygote`](@ref).
+  - [`AutoEnzyme`](@ref): Uses `Enzyme.jl` Reverse Mode and should be considered
+    experimental.
+
+!!! note
+
+    If `PolyesterForwardDiff.jl` is installed and loaded, then `SimpleNonlinearSolve.jl`
+    will automatically use `AutoPolyesterForwardDiff` as the default AD backend.
+
+!!! note
+
+    The `Sparse` versions of the methods refers to automated sparsity detection. These
+    methods automatically discover the sparse Jacobian form from the function `f`. Note that
+    all methods specialize the differentiation on a sparse Jacobian if the sparse Jacobian
+    is given as `prob.f.jac_prototype` in the `NonlinearFunction` definition, and the
+    `AutoSparse` here simply refers to whether this `jac_prototype` should be generated
+    automatically. For more details, see
+    [SparseDiffTools.jl](https://github.com/JuliaDiff/SparseDiffTools.jl) and
+    [Sparsity Detection Manual Entry](@ref sparsity-detection).
+
+## API Reference
+
+### Finite Differencing Backends
+
+```@docs
+AutoFiniteDiff
+AutoSparseFiniteDiff
+```
+
+### Forward Mode AD Backends
+
+```@docs
+AutoForwardDiff
+AutoSparseForwardDiff
+AutoPolyesterForwardDiff
+```
+
+### Reverse Mode AD Backends
+
+```@docs
+AutoZygote
+AutoSparseZygote
+AutoEnzyme
+NonlinearSolve.AutoSparseEnzyme
+```

src/NonlinearSolve.jl

@@ -24,7 +24,7 @@ import PrecompileTools: @recompile_invalidations, @compile_workload, @setup_work
 
     import SciMLBase: AbstractNonlinearAlgorithm, JacobianWrapper, AbstractNonlinearProblem,
         AbstractSciMLOperator, NLStats, _unwrap_val, has_jac, isinplace
-    import SparseDiffTools: AbstractSparsityDetection
+    import SparseDiffTools: AbstractSparsityDetection, AutoSparseEnzyme
     import StaticArraysCore: StaticArray, SVector, SArray, MArray, Size, SMatrix, MMatrix
 end
 
@@ -40,6 +40,7 @@ const True = Val(true)
 const False = Val(false)
 
 include("abstract_types.jl")
+include("adtypes.jl")
 include("timer_outputs.jl")
 include("internal/helpers.jl")
 

src/adtypes.jl

@@ -0,0 +1,148 @@
+# This just documents the AD types from ADTypes.jl
+
+"""
+    AutoFiniteDiff(; fdtype = Val(:forward), fdjtype = fdtype, fdhtype = Val(:hcentral))
+
+This uses [FiniteDiff.jl](https://github.com/JuliaDiff/FiniteDiff.jl). While not necessarily
+the most efficient, this is the only choice that doesn't require the `f` function to be
+automatically differentiable, which means it applies to any choice. However, because it's
+using finite differencing, one needs to be careful as this procedure introduces numerical
+error into the derivative estimates.
+
+  - Compatible with GPUs
+  - Can be used for Jacobian-Vector Products (JVPs)
+  - Can be used for Vector-Jacobian Products (VJPs)
+  - Supports both inplace and out-of-place functions
+
+### Keyword Arguments
+
+  - `fdtype`: the method used for defining the gradient
+  - `fdjtype`: the method used for defining the Jacobian of constraints.
+  - `fdhtype`: the method used for defining the Hessian
+"""
+AutoFiniteDiff
+
+"""
+    AutoSparseFiniteDiff()
+
+Sparse Version of [`AutoFiniteDiff`](@ref) that uses
+[FiniteDiff.jl](https://github.com/JuliaDiff/FiniteDiff.jl) and the column color vector of
+the Jacobian Matrix to efficiently compute the Sparse Jacobian.
+
+  - Supports both inplace and out-of-place functions
+"""
+AutoSparseFiniteDiff
+
+"""
+    AutoForwardDiff(; chunksize = nothing, tag = nothing)
+    AutoForwardDiff{chunksize, tagType}(tag::tagType)
+
+This uses the [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl) package. It is
+the fastest choice for square or wide systems. It is easy to use and compatible with most
+Julia functions which have loose type restrictions.
+
+  - Compatible with GPUs
+  - Can be used for Jacobian-Vector Products (JVPs)
+  - Supports both inplace and out-of-place functions
+
+For type-stability of internal operations, a positive `chunksize` must be provided.
+
+### Keyword Arguments
+
+  - `chunksize`: Count of dual numbers that can be propagated simultaneously. Setting this
+    number to a high value will lead to slowdowns. Use
+    [`NonlinearSolve.pickchunksize`](@ref) to get a proper value.
+  - `tag`: Used to avoid perturbation confusion. If set to `nothing`, we use a custom tag.
+"""
+AutoForwardDiff
+
+"""
+    AutoSparseForwardDiff(; chunksize = nothing, tag = nothing)
+    AutoSparseForwardDiff{chunksize, tagType}(tag::tagType)
+
+Sparse Version of [`AutoForwardDiff`](@ref) that uses
+[ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl) and the column color vector of
+the Jacobian Matrix to efficiently compute the Sparse Jacobian.
+
+  - Supports both inplace and out-of-place functions
+
+For type-stability of internal operations, a positive `chunksize` must be provided.
+
+### Keyword Arguments
+
+  - `chunksize`: Count of dual numbers that can be propagated simultaneously. Setting this
+    number to a high value will lead to slowdowns. Use
+    [`NonlinearSolve.pickchunksize`](@ref) to get a proper value.
+  - `tag`: Used to avoid perturbation confusion. If set to `nothing`, we use a custom tag.
+"""
+AutoSparseForwardDiff
+
+"""
+    AutoPolyesterForwardDiff(; chunksize = nothing)
+
+Uses [`PolyesterForwardDiff.jl`](https://github.com/JuliaDiff/PolyesterForwardDiff.jl)
+to compute the jacobian. This is essentially parallelized `ForwardDiff.jl`.
+
+    - Supports both inplace and out-of-place functions
+
+### Keyword Arguments
+
+    - `chunksize`: Count of dual numbers that can be propagated simultaneously. Setting
+      this number to a high value will lead to slowdowns. Use
+      [`NonlinearSolve.pickchunksize`](@ref) to get a proper value.
+"""
+AutoPolyesterForwardDiff
+
+"""
+    AutoZygote()
+
+Uses [`Zygote.jl`](https://github.com/FluxML/Zygote.jl) package. This is the staple
+reverse-mode AD that handles a large portion of Julia with good efficiency.
+
+  - Compatible with GPUs
+  - Can be used for Vector-Jacobian Products (VJPs)
+  - Supports only out-of-place functions
+
+For VJPs this is the current best choice. This is the most efficient method for long
+jacobians.
+"""
+AutoZygote
+
+"""
+    AutoSparseZygote()
+
+Sparse version of [`AutoZygote`](@ref) that uses
+[`Zygote.jl`](https://github.com/FluxML/Zygote.jl) and the row color vector of
+the Jacobian Matrix to efficiently compute the Sparse Jacobian.
+
+  - Supports only out-of-place functions
+
+This is efficient only for long jacobians or if the maximum value of the row color vector is
+significantly lower than the maximum value of the column color vector.
+"""
+AutoSparseZygote
+
+"""
+    AutoEnzyme()
+
+Uses reverse mode [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl). This is currently
+experimental, and not extensively tested on our end. We only support Jacobian construction
+and VJP support is currently not implemented.
+
+  - Supports both inplace and out-of-place functions
+"""
+AutoEnzyme
+
+"""
+    AutoSparseEnzyme()
+
+Sparse version of [`AutoEnzyme`](@ref) that uses
+[Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl) and the row color vector of
+the Jacobian Matrix to efficiently compute the Sparse Jacobian.
+
+  - Supports both inplace and out-of-place functions
+
+This is efficient only for long jacobians or if the maximum value of the row color vector is
+significantly lower than the maximum value of the column color vector.
+"""
+AutoSparseEnzyme