Skip to content
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.

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

Upgrade to ADTypes v1 #426

Merged
merged 10 commits into from
May 16, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
16 changes: 8 additions & 8 deletions Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ NonlinearSolveSymbolicsExt = "Symbolics"
NonlinearSolveZygoteExt = "Zygote"

[compat]
ADTypes = "0.2.6"
ADTypes = "1.1.0"
Aqua = "0.8"
ArrayInterface = "7.9"
BandedMatrices = "1.5"
Expand All @@ -69,7 +69,7 @@ Enzyme = "0.11.15, 0.12"
FastBroadcast = "0.2.8"
FastClosures = "0.3.2"
FastLevenbergMarquardt = "0.1"
FiniteDiff = "2.21"
FiniteDiff = "2.22"
FixedPointAcceleration = "0.3"
ForwardDiff = "0.10.36"
LazyArrays = "1.8.2"
Expand All @@ -83,7 +83,7 @@ NLSolvers = "0.5"
NLsolve = "4.5"
NaNMath = "1"
NonlinearProblemLibrary = "0.1.2"
OrdinaryDiffEq = "6.74"
OrdinaryDiffEq = "6.75"
Pkg = "1.10"
PrecompileTools = "1.2"
Preferences = "1.4"
Expand All @@ -94,17 +94,17 @@ RecursiveArrayTools = "3.8"
Reexport = "1.2"
SIAMFANLEquations = "1.0.1"
SafeTestsets = "0.1"
SciMLBase = "2.28.0"
SimpleNonlinearSolve = "1.2"
SciMLBase = "2.34.0"
SimpleNonlinearSolve = "1.8"
SparseArrays = "1.10"
SparseDiffTools = "2.17"
SparseDiffTools = "2.19"
SpeedMapping = "0.3"
StableRNGs = "1"
StaticArrays = "1.7"
StaticArrays = "1.9"
StaticArraysCore = "1.4"
Sundials = "4.23.1"
Symbolics = "5.13"
SymbolicIndexingInterface = "0.3.3"
SymbolicIndexingInterface = "0.3.15"
Test = "1.10"
TimerOutputs = "0.5.23"
Zygote = "0.6.69"
Expand Down
16 changes: 7 additions & 9 deletions docs/src/basics/autodiff.md
Original file line number Diff line number Diff line change
Expand Up @@ -3,19 +3,16 @@
## 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.

Expand All @@ -26,37 +23,38 @@

!!! note

The `Sparse` versions of the methods refers to automated sparsity detection. These
The sparse versions of the methods refer 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).
[Sparsity Detection Manual Entry](@ref sparsity-detection), as well as the
documentation of [ADTypes.jl](https://github.com/SciML/ADTypes.jl).

## API Reference

```@docs
AutoSparse
```

### Finite Differencing Backends

```@docs
AutoFiniteDiff
AutoSparseFiniteDiff
```

### Forward Mode AD Backends

```@docs
AutoForwardDiff
AutoSparseForwardDiff
AutoPolyesterForwardDiff
```

### Reverse Mode AD Backends

```@docs
AutoZygote
AutoSparseZygote
AutoEnzyme
NonlinearSolve.AutoSparseEnzyme
```
4 changes: 2 additions & 2 deletions docs/src/basics/sparsity_detection.md
Original file line number Diff line number Diff line change
Expand Up @@ -59,9 +59,9 @@ refer to the documentation there for more details.
If you constructed a Nonlinear Solver with a sparse AD type, for example

```julia
NewtonRaphson(; autodiff = AutoSparseForwardDiff())
NewtonRaphson(; autodiff = AutoSparse(AutoForwardDiff()))
# OR
TrustRegion(; autodiff = AutoSparseZygote())
TrustRegion(; autodiff = AutoSparse(AutoZygote()))
```

then NonlinearSolve will automatically perform matrix coloring and use sparse
Expand Down
12 changes: 6 additions & 6 deletions docs/src/tutorials/large_systems.md
Original file line number Diff line number Diff line change
Expand Up @@ -128,7 +128,7 @@ include:
In the next section, we will discuss how to declare a sparse Jacobian and how to use
[Symbolics.jl](https://github.com/JuliaSymbolics/Symbolics.jl), to compute exact sparse
jacobians. This is triggered if you pass in a sparse autodiff type such as
`AutoSparseForwardDiff()`. If `Symbolics.jl` is loaded, then the default changes to
`AutoSparse(AutoForwardDiff())`. If `Symbolics.jl` is loaded, then the default changes to
Symbolic Sparsity Detection. See the manual entry on
[Sparsity Detection](@ref sparsity-detection) for more details on the default.

Expand All @@ -137,13 +137,13 @@ using BenchmarkTools # for @btime
@btime solve(prob_brusselator_2d, NewtonRaphson());
@btime solve(prob_brusselator_2d,
NewtonRaphson(; autodiff = AutoSparseForwardDiff(; chunksize = 32)));
NewtonRaphson(; autodiff = AutoSparse(AutoForwardDiff(; chunksize = 32))));
@btime solve(prob_brusselator_2d,
NewtonRaphson(;
autodiff = AutoSparseForwardDiff(; chunksize = 32), linsolve = KLUFactorization()));
NewtonRaphson(; autodiff = AutoSparse(AutoForwardDiff(; chunksize = 32)),
linsolve = KLUFactorization()));
@btime solve(prob_brusselator_2d,
NewtonRaphson(;
autodiff = AutoSparseForwardDiff(; chunksize = 32), linsolve = KrylovJL_GMRES()));
NewtonRaphson(; autodiff = AutoSparse(AutoForwardDiff(; chunksize = 32)),
linsolve = KrylovJL_GMRES()));
nothing # hide
```

Expand Down
29 changes: 13 additions & 16 deletions src/NonlinearSolve.jl
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ import PrecompileTools: @recompile_invalidations, @compile_workload, @setup_work

import SciMLBase: AbstractNonlinearAlgorithm, JacobianWrapper, AbstractNonlinearProblem,
AbstractSciMLOperator, NLStats, _unwrap_val, has_jac, isinplace
import SparseDiffTools: AbstractSparsityDetection, AutoSparseEnzyme
import SparseDiffTools: AbstractSparsityDetection
import StaticArraysCore: StaticArray, SVector, SArray, MArray, Size, SMatrix, MMatrix
import SymbolicIndexingInterface: SymbolicIndexingInterface, ParameterIndexingProxy,
symbolic_container, parameter_values, state_values,
Expand All @@ -36,9 +36,6 @@ end

@reexport using ADTypes, SciMLBase, SimpleNonlinearSolve

const AbstractSparseADType = Union{ADTypes.AbstractSparseFiniteDifferences,
ADTypes.AbstractSparseForwardMode, ADTypes.AbstractSparseReverseMode}

# Type-Inference Friendly Check for Extension Loading
is_extension_loaded(::Val) = false

Expand Down Expand Up @@ -121,18 +118,18 @@ include("default.jl")

@compile_workload begin
@sync begin
for T in (Float32, Float64), (fn, u0) in nlfuncs
Threads.@spawn NonlinearProblem(fn, T.(u0), T(2))
end
for (fn, u0) in nlfuncs
Threads.@spawn NonlinearLeastSquaresProblem(fn, u0, 2.0)
end
for prob in probs_nls, alg in nls_algs
Threads.@spawn solve(prob, alg; abstol = 1e-2, verbose = false)
end
for prob in probs_nlls, alg in nlls_algs
Threads.@spawn solve(prob, alg; abstol = 1e-2, verbose = false)
end
for T in (Float32, Float64), (fn, u0) in nlfuncs
Threads.@spawn NonlinearProblem(fn, T.(u0), T(2))
end
for (fn, u0) in nlfuncs
Threads.@spawn NonlinearLeastSquaresProblem(fn, u0, 2.0)
end
for prob in probs_nls, alg in nls_algs
Threads.@spawn solve(prob, alg; abstol = 1e-2, verbose = false)
end
for prob in probs_nlls, alg in nlls_algs
Threads.@spawn solve(prob, alg; abstol = 1e-2, verbose = false)
end
end
end
end
Expand Down
88 changes: 40 additions & 48 deletions src/adtypes.jl
Original file line number Diff line number Diff line change
Expand Up @@ -22,17 +22,6 @@ error into the derivative estimates.
"""
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)
Expand All @@ -56,27 +45,6 @@ For type-stability of internal operations, a positive `chunksize` must be provid
"""
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)
Expand Down Expand Up @@ -108,20 +76,6 @@ 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()
Expand All @@ -134,15 +88,53 @@ and VJP support is currently not implemented.
AutoEnzyme

"""
AutoSparseEnzyme()
AutoSparse(AutoEnzyme())
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.
AutoSparse(AutoFiniteDiff())
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
AutoSparse(AutoForwardDiff(; chunksize = nothing, tag = nothing))
AutoSparse(AutoForwardDiff{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.
AutoSparse(AutoZygote())
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.
"""
AutoSparseEnzyme
AutoSparse
7 changes: 4 additions & 3 deletions src/algorithms/trust_region.jl
Original file line number Diff line number Diff line change
Expand Up @@ -26,13 +26,14 @@ function TrustRegion(; concrete_jac = nothing, linsolve = nothing, precs = DEFAU
shrink_factor::Real = 1 // 4, expand_factor::Real = 2 // 1,
max_shrink_times::Int = 32, autodiff = nothing, vjp_autodiff = nothing)
descent = Dogleg(; linsolve, precs)
if autodiff isa
Union{ADTypes.AbstractForwardMode, ADTypes.AbstractFiniteDifferencesMode}
if autodiff !== nothing && ADTypes.mode(autodiff) isa ADTypes.ForwardMode
forward_ad = autodiff
else
forward_ad = nothing
end
if isnothing(vjp_autodiff) && autodiff isa ADTypes.AbstractFiniteDifferencesMode
if isnothing(vjp_autodiff) &&
autodiff isa Union{ADTypes.AutoFiniteDiff, ADTypes.AutoFiniteDifferences}
# TODO: why not just ForwardMode?
gdalle marked this conversation as resolved.
Show resolved Hide resolved
vjp_autodiff = autodiff
end
trustregion = GenericTrustRegionScheme(;
Expand Down
12 changes: 8 additions & 4 deletions src/core/generalized_first_order.jl
Original file line number Diff line number Diff line change
Expand Up @@ -55,10 +55,14 @@ function GeneralizedFirstOrderAlgorithm{concrete_jac, name}(;
descent, linesearch = missing, trustregion = missing,
jacobian_ad = nothing, forward_ad = nothing, reverse_ad = nothing,
max_shrink_times::Int = typemax(Int)) where {concrete_jac, name}
forward_ad = ifelse(forward_ad !== nothing, forward_ad,
ifelse(jacobian_ad isa ADTypes.AbstractForwardMode, jacobian_ad, nothing))
reverse_ad = ifelse(reverse_ad !== nothing, reverse_ad,
ifelse(jacobian_ad isa ADTypes.AbstractReverseMode, jacobian_ad, nothing))
forward_ad = ifelse(forward_ad !== nothing,
forward_ad,
ifelse(jacobian_ad !== nothing && ADTypes.mode(jacobian_ad) isa ADTypes.ForwardMode,
jacobian_ad, nothing))
reverse_ad = ifelse(reverse_ad !== nothing,
reverse_ad,
ifelse(jacobian_ad !== nothing && ADTypes.mode(jacobian_ad) isa ADTypes.ReverseMode,
jacobian_ad, nothing))

if linesearch !== missing && !(linesearch isa AbstractNonlinearSolveLineSearchAlgorithm)
Base.depwarn("Passing in a `LineSearches.jl` algorithm directly is deprecated. \
Expand Down
2 changes: 1 addition & 1 deletion src/globalization/line_search.jl
Original file line number Diff line number Diff line change
Expand Up @@ -122,7 +122,7 @@ function __internal_init(
end
else
autodiff = get_concrete_reverse_ad(
alg.autodiff, prob; check_forward_mode = true)
alg.autodiff, prob; check_reverse_mode = true)
vjp_op = VecJacOperator(prob, fu, u; autodiff)
if isinplace(prob)
g_cache = similar(u)
Expand Down
2 changes: 1 addition & 1 deletion src/internal/approximate_initialization.jl
Original file line number Diff line number Diff line change
Expand Up @@ -149,7 +149,7 @@ function __internal_init(
prob::AbstractNonlinearProblem, alg::TrueJacobianInitialization, solver, f::F, fu,
u, p; linsolve = missing, internalnorm::IN = DEFAULT_NORM, kwargs...) where {F, IN}
autodiff = get_concrete_forward_ad(
alg.autodiff, prob; check_reverse_mode = false, kwargs...)
alg.autodiff, prob; check_forward_mode = false, kwargs...)
jac_cache = JacobianCache(prob, solver, prob.f, fu, u, p; autodiff, linsolve)
J = alg.structure(jac_cache(nothing))
return InitializedApproximateJacobianCache(
Expand Down
Loading
Loading