Restructure sparse backends and replace subtyping by traits (#40)

Author: gdalle

.github/workflows/CI.yml

@@ -18,9 +18,8 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.6'
+          - '1.10'
           - '1'
-          - 'nightly'
         os:
           - ubuntu-latest
         arch:
@@ -34,3 +33,9 @@ jobs:
       - uses: julia-actions/cache@v1
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
+      - uses: julia-actions/julia-processcoverage@v1
+      - uses: codecov/codecov-action@v4
+        with:
+          files: lcov.info
+          token: ${{ secrets.CODECOV_TOKEN }}
+          fail_ci_if_error: true

Project.toml

@@ -1,13 +1,29 @@
 name = "ADTypes"
 uuid = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
-authors = ["Vaibhav Dixit <[email protected]> and contributors"]
-version = "0.2.7"
+authors = [
+    "Vaibhav Dixit <[email protected]>, Guillaume Dalle and contributors",
+]
+version = "1.0.0"
+
+[weakdeps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
+
+[extensions]
+ADTypesChainRulesCoreExt = "ChainRulesCore"
+ADTypesEnzymeCoreExt = "EnzymeCore"
 
 [compat]
-julia = "1.6"
+ChainRulesCore = "1.23.0"
+EnzymeCore = "0.7.2"
+julia = "1.10"
 
 [extras]
+Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test"]
+test = ["Aqua", "ChainRulesCore", "EnzymeCore", "JET", "Test"]

README.md

@@ -4,29 +4,24 @@
 [![Docs dev](https://img.shields.io/badge/docs-dev-blue.svg)](https://SciML.github.io/ADTypes.jl/dev/)
 [![Build Status](https://github.com/SciML/ADTypes.jl/actions/workflows/CI.yml/badge.svg?branch=main)](https://github.com/SciML/ADTypes.jl/actions/workflows/CI.yml?query=branch%3Amain)
 
-ADTypes.jl is a multi-valued logic system specifying the choice of an automatic differentiation (AD) library and its parameters.
+ADTypes.jl is a multi-valued logic system to choose an automatic differentiation (AD) package and specify its parameters.
 
 ## Which AD libraries are supported?
 
 See the API reference in the documentation.
+If a given package is missing, feel free to open an issue or pull request.
 
-## Why should packages adopt this standard?
+## Why should AD users adopt this standard?
 
-A common practice is the use of a boolean keyword argument like `autodiff = true/false`.
-However, boolean logic is not precise enough for all the choices required.
-For instance, forward mode AD is implemented by both ForwardDiff and Enzyme, which makes `autodiff = true` ambiguous.
-Something like `ChooseForwardDiff()` is thus required, possibly with additional parameters depending on the library.
+A natural approach is to use a keyword argument with e.g. `Bool` or `Symbol` values.
+Let's see a few examples to understand why this is not enough:
 
-The risk is that every package developer might develop their own version of `ChooseForwardDiff()`, which would ruin interoperability.
-This is why ADTypes.jl provides a single set of shared types for this task, as an extremely lightweight dependency.
-Wonder no more: `ADTypes.AutoForwardDiff()` is the way to go.
-
-## Why define types instead of enums?
-
-If we used enums, they would not contain type-level information useful for dispatch.
-This is needed by many AD libraries to ensure type stability.
-Notably, the choice of config or cache type is different with each AD, so we must know statically which AD library is chosen.
+  - `autodiff = true`: ambiguous, we don't know which AD package should be used
+  - `autodiff = :forward`: ambiguous, there are several AD packages implementing both forward and reverse mode (and there are other modes beyond that)
+  - `autodiff = :Enzyme`: ambiguous, some AD packages can work both in forward and reverse mode
+  - `autodiff = (:Enzyme, :forward)`: not too bad, but many AD packages require additional configuration (number of chunks, tape compilation, etc.)
 
-## Why is this AD package missing?
-
-Feel free to open a pull request adding it.
+A more involved struct is thus required, with package-specific parameters.
+If every AD user develops their own version of said struct, it will ruin interoperability.
+This is why ADTypes.jl provides a single set of shared types for this task, as an extremely lightweight dependency.
+They are types and not enums because we need AD choice information statically to use it for dispatch.

docs/src/index.md

@@ -7,16 +7,87 @@ CollapsedDocStrings = true
 
 Documentation for [ADTypes.jl](https://github.com/SciML/ADTypes.jl/).
 
-## Public
+```@docs
+ADTypes
+AbstractADType
+```
+
+## Dense AD
+
+### Forward mode
+
+Algorithmic differentiation:
+
+```@docs
+AutoForwardDiff
+AutoPolyesterForwardDiff
+```
+
+Finite differences:
+
+```@docs
+AutoFiniteDiff
+AutoFiniteDifferences
+```
+
+### Reverse mode
+
+```@docs
+AutoReverseDiff
+AutoTapir
+AutoTracker
+AutoZygote
+```
+
+### Forward or reverse mode
+
+```@docs
+AutoEnzyme
+AutoChainRules
+AutoDiffractor
+```
+
+### Symbolic mode
+
+```@docs
+AutoFastDifferentiation
+AutoSymbolics
+```
+
+## Sparse AD
+
+```@docs
+AutoSparse
+ADTypes.dense_ad
+```
+
+### Sparsity detector
+
+```@docs
+ADTypes.sparsity_detector
+ADTypes.AbstractSparsityDetector
+ADTypes.jacobian_sparsity
+ADTypes.hessian_sparsity
+ADTypes.NoSparsityDetector
+```
+
+### Coloring algorithm
 
-```@autodocs
-Modules = [ADTypes]
-Private = false
+```@docs
+ADTypes.coloring_algorithm
+ADTypes.AbstractColoringAlgorithm
+ADTypes.column_coloring
+ADTypes.row_coloring
+ADTypes.NoColoringAlgorithm
 ```
 
-## Internal
+## Modes
 
-```@autodocs
-Modules = [ADTypes]
-Public = false
+```@docs
+ADTypes.mode
+ADTypes.AbstractMode
+ADTypes.ForwardMode
+ADTypes.ForwardOrReverseMode
+ADTypes.ReverseMode
+ADTypes.SymbolicMode
 ```

ext/ADTypesChainRulesCoreExt.jl

@@ -0,0 +1,29 @@
+module ADTypesChainRulesCoreExt
+
+using ADTypes: ADTypes, AutoChainRules
+using ChainRulesCore: HasForwardsMode, HasReverseMode,
+                      NoForwardsMode, NoReverseMode,
+                      RuleConfig
+
+# see https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/superpowers/ruleconfig.html
+
+function ADTypes.mode(::AutoChainRules{RC}) where {
+        RC <: RuleConfig{>:HasForwardsMode}
+}
+    return ADTypes.ForwardMode()
+end
+
+function ADTypes.mode(::AutoChainRules{RC}) where {
+        RC <: RuleConfig{>:HasReverseMode}
+}
+    return ADTypes.ReverseMode()
+end
+
+function ADTypes.mode(::AutoChainRules{RC}) where {
+        RC <: RuleConfig{>:Union{HasForwardsMode, HasReverseMode}}
+}
+    # more specific than the previous two
+    return ADTypes.ForwardOrReverseMode()
+end
+
+end

ext/ADTypesEnzymeCoreExt.jl

@@ -0,0 +1,9 @@
+module ADTypesEnzymeCoreExt
+
+using ADTypes: ADTypes, AutoEnzyme
+using EnzymeCore: EnzymeCore
+
+ADTypes.mode(::AutoEnzyme{M}) where {M <: EnzymeCore.ForwardMode} = ADTypes.ForwardMode()
+ADTypes.mode(::AutoEnzyme{M}) where {M <: EnzymeCore.ReverseMode} = ADTypes.ReverseMode()
+
+end