Upgrade to ADTypes v1.0 (#194)

* Upgrade to ADTypes v1.0 * Fix docs * Fix docs * Aqua * Fix tests * Fix CI * Revert version bump for now * Add manifests * Fix compat * Fix tuto * Re-add Symbolics * Add coloring algorithms * Fix CI * Rev * Resolve * Remove downloads in favor of manifests * Don't force latest version * Sparsity works * Deactivate workflows * SciMLBase update * Remove playground * Reactivate workflows * Rely on registered versions * Remove warnonly in docs * Push preview * Rm SciMLBase
JuliaDiff · Apr 22, 2024 · 1d39604 · 1d39604
1 parent 4fa52a2
commit 1d39604
Show file tree

Hide file tree

Showing 52 changed files with 784 additions and 578 deletions.
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
@@ -56,7 +56,6 @@ jobs:
       - uses: julia-actions/julia-runtest@v1
         with:
           project: ${{ matrix.pkg.dir }}
-          force_latest_compatible_version: true
       - uses: julia-actions/julia-processcoverage@v1
         with:
           directories: ${{ matrix.pkg.dir }}/src,${{ matrix.pkg.dir }}/ext

diff --git a/DifferentiationInterface/Project.toml b/DifferentiationInterface/Project.toml
@@ -1,7 +1,7 @@
 name = "DifferentiationInterface"
 uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.2.1"
+version = "0.3.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -35,14 +35,13 @@ DifferentiationInterfaceFiniteDifferencesExt = "FiniteDifferences"
 DifferentiationInterfaceForwardDiffExt = "ForwardDiff"
 DifferentiationInterfacePolyesterForwardDiffExt = "PolyesterForwardDiff"
 DifferentiationInterfaceReverseDiffExt = "ReverseDiff"
-DifferentiationInterfaceSparseDiffToolsExt = ["SparseDiffTools", "Symbolics"]
 DifferentiationInterfaceSymbolicsExt = "Symbolics"
 DifferentiationInterfaceTapirExt = "Tapir"
 DifferentiationInterfaceTrackerExt = "Tracker"
 DifferentiationInterfaceZygoteExt = "Zygote"
 
 [compat]
-ADTypes = "0.2.7"
+ADTypes = "1.0.0"
 ChainRulesCore = "1.23.0"
 Diffractor = "=0.2.6"
 DocStringExtensions = "0.9.3"
@@ -55,7 +54,6 @@ ForwardDiff = "0.10.36"
 LinearAlgebra = "1"
 PolyesterForwardDiff = "0.1.1"
 ReverseDiff = "1.15.1"
-SparseDiffTools = "2.17.0"
 Symbolics = "5.27.1"
 Tapir = "0.1.2"
 Test = "1"
@@ -79,34 +77,11 @@ Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
-SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Tapir = "07d77754-e150-4737-8c94-cd238a1fb45b"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = [
-    "ADTypes",
-    "Aqua",
-    "Diffractor",
-    "Documenter",
-    "Enzyme",
-    "FastDifferentiation",
-    "FiniteDiff",
-    "FiniteDifferences",
-    "ForwardDiff",
-    "JET",
-    "JuliaFormatter",
-    "Pkg",
-    "PolyesterForwardDiff",
-    "ReverseDiff",
-    "SparseArrays",
-    "SparseDiffTools",
-    "Symbolics",
-    "Tapir",
-    "Test",
-    "Tracker",
-    "Zygote",
-]
+test = ["ADTypes", "Aqua", "Diffractor", "Documenter", "Enzyme", "FastDifferentiation", "FiniteDiff", "FiniteDifferences", "ForwardDiff", "JET", "JuliaFormatter", "Pkg", "PolyesterForwardDiff", "ReverseDiff", "SparseArrays", "Symbolics", "Tapir", "Test", "Tracker", "Zygote"]
diff --git a/DifferentiationInterface/README.md b/DifferentiationInterface/README.md
@@ -28,29 +28,22 @@ This package provides a backend-agnostic syntax to differentiate functions of th
 
 ## Compatibility
 
-We support most of the backends defined by [ADTypes.jl](https://github.com/SciML/ADTypes.jl):
-
-| Backend                                                                         | Object                                                                   |
-| :------------------------------------------------------------------------------ | :----------------------------------------------------------------------- |
-| [ChainRulesCore.jl](https://github.com/JuliaDiff/ChainRulesCore.jl)             | `AutoChainRules(; ruleconfig)`                                           |
-| [Diffractor.jl](https://github.com/JuliaDiff/Diffractor.jl)                     | `AutoDiffractor()`                                                       |
-| [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl)                              | `AutoEnzyme(; mode=Enzyme.Forward)`, `AutoEnzyme(; mode=Enzyme.Reverse)` |
-| [FiniteDiff.jl](https://github.com/JuliaDiff/FiniteDiff.jl)                     | `AutoFiniteDiff()`                                                       |
-| [FiniteDifferences.jl](https://github.com/JuliaDiff/FiniteDifferences.jl)       | `AutoFiniteDifferences(; fdm)`                                           |
-| [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl)                   | `AutoForwardDiff()`                                                      |
-| [PolyesterForwardDiff.jl](https://github.com/JuliaDiff/PolyesterForwardDiff.jl) | `AutoPolyesterForwardDiff(; chunksize)`                                  |
-| [ReverseDiff.jl](https://github.com/JuliaDiff/ReverseDiff.jl)                   | `AutoReverseDiff()`                                                      |
-| [SparseDiffTools.jl](https://github.com/JuliaDiff/SparseDiffTools.jl)           | `AutoSparseForwardDiff()`, `AutoSparseFiniteDiff()`                      |
-| [Tracker.jl](https://github.com/FluxML/Tracker.jl)                              | `AutoTracker()`                                                          |
-| [Zygote.jl](https://github.com/FluxML/Zygote.jl)                                | `AutoZygote()`                                                           |
-
-We also provide some experimental backends ourselves:
-
-| Backend                                                                          | Object                                                         |
-| :------------------------------------------------------------------------------- | :------------------------------------------------------------- |
-| [FastDifferentiation.jl](https://github.com/brianguenter/FastDifferentiation.jl) | `AutoFastDifferentiation()`, `AutoSparseFastDifferentiation()` |
-| [Symbolics.jl](https://github.com/JuliaSymbolics/Symbolics.jl)                   | `AutoSymbolics()`, `AutoSparseSymbolics()`                     |
-| [Tapir.jl](https://github.com/withbayes/Tapir.jl)                                | `AutoTapir()`                                                  |
+We support all of the backends defined by [ADTypes.jl](https://github.com/SciML/ADTypes.jl) v1.0:
+
+- [ChainRulesCore.jl](https://github.com/JuliaDiff/ChainRulesCore.jl)
+- [Diffractor.jl](https://github.com/JuliaDiff/Diffractor.jl)
+- [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl)
+- [FastDifferentiation.jl](https://github.com/brianguenter/FastDifferentiation.jl)
+- [FiniteDiff.jl](https://github.com/JuliaDiff/FiniteDiff.jl)
+- [FiniteDifferences.jl](https://github.com/JuliaDiff/FiniteDifferences.jl)
+- [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl)
+- [PolyesterForwardDiff.jl](https://github.com/JuliaDiff/PolyesterForwardDiff.jl)
+- [ReverseDiff.jl](https://github.com/JuliaDiff/ReverseDiff.jl)
+- [SparseDiffTools.jl](https://github.com/JuliaDiff/SparseDiffTools.jl)
+- [Symbolics.jl](https://github.com/JuliaSymbolics/Symbolics.jl) (currently broken)
+- [Tapir.jl](https://github.com/withbayes/Tapir.jl)
+- [Tracker.jl](https://github.com/FluxML/Tracker.jl)
+- [Zygote.jl](https://github.com/FluxML/Zygote.jl)
 
 ## Installation
 
@@ -76,11 +69,11 @@ julia> Pkg.add(
 ## Example
 
 ```jldoctest readme
-julia> import ForwardDiff
+julia> import ADTypes, ForwardDiff
 
 julia> using DifferentiationInterface
 
-julia> backend = AutoForwardDiff();
+julia> backend = ADTypes.AutoForwardDiff();
 
 julia> f(x) = sum(abs2, x);
 

diff --git a/DifferentiationInterface/docs/Project.toml b/DifferentiationInterface/docs/Project.toml
@@ -16,12 +16,11 @@ Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
 PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
-SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Tapir = "07d77754-e150-4737-8c94-cd238a1fb45b"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
-ADTypes = "0.2.7"
+DifferentiationInterface = "0.3"
 Documenter = "1"
diff --git a/DifferentiationInterface/docs/make.jl b/DifferentiationInterface/docs/make.jl
@@ -14,7 +14,6 @@ using ForwardDiff: ForwardDiff
 using PolyesterForwardDiff: PolyesterForwardDiff
 using ReverseDiff: ReverseDiff
 using Symbolics: Symbolics
-using SparseDiffTools: SparseDiffTools
 using Tapir: Tapir
 using Tracker: Tracker
 using Zygote: Zygote
@@ -31,14 +30,14 @@ extensions = [
     get_extension(DI, :DifferentiationInterfaceForwardDiffExt),
     get_extension(DI, :DifferentiationInterfacePolyesterForwardDiffExt),
     get_extension(DI, :DifferentiationInterfaceReverseDiffExt),
-    get_extension(DI, :DifferentiationInterfaceSparseDiffToolsExt),
+    get_extension(DI, :DifferentiationInterfaceSymbolicsExt),
     get_extension(DI, :DifferentiationInterfaceTapirExt),
     get_extension(DI, :DifferentiationInterfaceTrackerExt),
     get_extension(DI, :DifferentiationInterfaceZygoteExt),
 ]
 
 makedocs(;
-    modules=[ADTypes, DifferentiationInterface, extensions...],
+    modules=[DifferentiationInterface, extensions...],
     authors="Guillaume Dalle, Adrian Hill",
     sitename="DifferentiationInterface.jl",
     format=Documenter.HTML(),
@@ -48,7 +47,6 @@ makedocs(;
         "API reference" => "api.md",
         "Advanced" => ["design.md", "extensions.md"],
     ],
-    warnonly=:missing_docs,  # missing docs for ADTypes.jl are normal
 )
 
 deploydocs(;

diff --git a/DifferentiationInterface/docs/src/backends.md b/DifferentiationInterface/docs/src/backends.md
@@ -4,6 +4,7 @@ CollapsedDocStrings = true
 ```
 
 ```@setup backends
+using ADTypes
 using DifferentiationInterface
 using DifferentiationInterface: backend_string
 import Markdown
@@ -12,72 +13,29 @@ import Diffractor, Enzyme, FastDifferentiation, FiniteDiff, FiniteDifferences, F
 function all_backends()
     return [
         AutoDiffractor(),
-        AutoEnzyme(Enzyme.Forward),
-        AutoEnzyme(Enzyme.Reverse),
+        AutoEnzyme(; mode=Enzyme.Forward),
+        AutoEnzyme(; mode=Enzyme.Reverse),
         AutoFastDifferentiation(),
         AutoFiniteDiff(),
-        AutoFiniteDifferences(FiniteDifferences.central_fdm(3, 1)),
+        AutoFiniteDifferences(; fdm=FiniteDifferences.central_fdm(3, 1)),
         AutoForwardDiff(),
         AutoPolyesterForwardDiff(; chunksize=1),
         AutoReverseDiff(),
-        AutoSymbolics(),
         AutoTapir(),
         AutoTracker(),
         AutoZygote(),
     ]
 end
-
-function all_backends_without_enzyme()
-    return filter(all_backends()) do b
-        !isa(b, AutoEnzyme)
-    end
-end
 ```
 
 # Backends
 
 ## Types
 
-We support (and re-export) most backend choices from [ADTypes.jl](https://github.com/SciML/ADTypes.jl), and we provide a few more of our own.
-
-!!! warning
-    Only the backends listed below are supported by DifferentiationInterface.jl, even though ADTypes.jl defines more.
-
-### Dense
-
-```@docs
-AutoChainRules
-AutoDiffractor
-AutoEnzyme
-AutoFastDifferentiation
-AutoForwardDiff
-AutoForwardDiff()
-AutoFiniteDiff
-AutoFiniteDifferences
-AutoPolyesterForwardDiff
-AutoPolyesterForwardDiff()
-AutoReverseDiff
-AutoSymbolics
-AutoTapir
-AutoTracker
-AutoZygote
-```
-
-### Sparse
+We support all dense backend choices from [ADTypes.jl](https://github.com/SciML/ADTypes.jl), as well as their sparse wrapper `AutoSparse`.
 
 For sparse backends, only the Jacobian and Hessian operators are implemented differently, the other operators behave the same as for the corresponding dense backend.
 
-```@docs
-AutoSparseFastDifferentiation
-AutoSparseFiniteDiff
-AutoSparseForwardDiff
-AutoSparseForwardDiff()
-AutoSparsePolyesterForwardDiff
-AutoSparseReverseDiff
-AutoSparseSymbolics
-AutoSparseZygote
-```
-
 ## Availability
 
 You can use [`check_available`](@ref) to verify whether a given backend is loaded, like we did below:

diff --git a/DifferentiationInterface/docs/src/design.md b/DifferentiationInterface/docs/src/design.md
@@ -5,17 +5,11 @@
 To be usable with DifferentiationInterface.jl, an AD backend needs an object subtyping `ADTypes.AbstractADType`.
 In addition, some operators must be defined:
 
-| backend subtype                               | pushforward necessary | pullback necessary |
-| :-------------------------------------------- | :-------------------- | :----------------- |
-| `ADTypes.AbstractForwardMode`                 | yes                   | no                 |
-| `ADTypes.AbstractFiniteDifferencesMode`       | yes                   | no                 |
-| `ADTypes.AbstractReverseMode`                 | no                    | yes                |
-| `ADTypes.AbstractSymbolicDifferentiationMode` | yes                   | yes                |
+| backend subtype        | pushforward necessary | pullback necessary |
+| :--------------------- | :-------------------- | :----------------- |
+| `ADTypes.ForwardMode`  | yes                   | no                 |
+| `ADTypes.ReverseMode`  | no                    | yes                |
+| `ADTypes.SymbolicMode` | yes                   | yes                |
 
 Every backend we support corresponds to a package extension of DifferentiationInterface.jl (located in the `ext` subfolder).
 Advanced users are welcome to code more backends and submit pull requests!
-
-## Fallback call structure
-
-!!! warning
-    This is still in flux, come back later!
diff --git a/DifferentiationInterface/docs/src/extensions.md b/DifferentiationInterface/docs/src/extensions.md
@@ -4,7 +4,7 @@
 CurrentModule = DifferentiationInterface
 ```
 
-Backend-specific extension content is not part of the public API.
+Package extension content is not part of the public API.
 If any docstrings are present in an extension, they will appear below.
 
 ```@autodocs
@@ -18,7 +18,7 @@ Modules = [
     Base.get_extension(DifferentiationInterface, :DifferentiationInterfaceForwardDiffExt),
     Base.get_extension(DifferentiationInterface, :DifferentiationInterfacePolyesterForwardDiffExt),
     Base.get_extension(DifferentiationInterface, :DifferentiationInterfaceReverseDiffExt),
-    Base.get_extension(DifferentiationInterface, :DifferentiationInterfaceSparseDiffToolsExt),
+    Base.get_extension(DifferentiationInterface, :DifferentiationInterfaceSymbolicsExt),
     Base.get_extension(DifferentiationInterface, :DifferentiationInterfaceTapirExt),
     Base.get_extension(DifferentiationInterface, :DifferentiationInterfaceTrackerExt),
     Base.get_extension(DifferentiationInterface, :DifferentiationInterfaceZygoteExt)

diff --git a/DifferentiationInterface/docs/src/overview.md b/DifferentiationInterface/docs/src/overview.md
@@ -109,16 +109,10 @@ We offer two ways to perform second-order differentiation (for [`second_derivati
 
 ### Sparsity
 
-[ADTypes.jl](https://github.com/SciML/ADTypes.jl) provides [sparse versions](@ref Sparse) of many common AD backends.
-They can accelerate the computation of sparse Jacobians and Hessians:
+[ADTypes.jl](https://github.com/SciML/ADTypes.jl) provides `AutoSparse` to accelerate the computation of sparse Jacobians and Hessians:
 
-- for sparse Jacobians, just select one of them as your first-order backend.
-- for sparse Hessians, select one of them as the _outer part_ of a [`SecondOrder`](@ref) backend (in that case, the Hessian is obtained as the sparse Jacobian of the gradient).
-
-The sparsity pattern is computed automatically with [Symbolics.jl](https://github.com/JuliaSymbolics/Symbolics.jl) during the preparation step.
-
-!!! info "Planned feature"
-    Modular sparsity pattern computation, with other algorithms beyond those from Symbolics.jl.
+- for sparse Jacobians, wrap `AutoSparse` around a first-order backend.
+- for sparse Hessians, wrap `AutoSparse` around a [`SecondOrder`](@ref) backend.
 
 ### Split reverse mode
 
@@ -129,12 +123,7 @@ We make this available for all backends with the following operators:
 | :--------------------------------- | :---------------------------------- |
 | [`value_and_pullback_split`](@ref) | [`value_and_pullback!_split`](@ref) |
 
-## Not supported
-
-### Batched evaluation
-
-!!! info "Planned feature"
-    Interface for providing several pushforward / pullback seeds at once, similar to the chunking in ForwardDiff.jl or the batches in Enzyme.jl.
+## Going further
 
 ### Non-standard types
 
@@ -147,3 +136,7 @@ We voluntarily keep the type annotations minimal, so that passing more complex o
 
 Restricting the API to one input and one output has many coding advantages, but it is not very flexible.
 If you need more than that, use [ComponentArrays.jl](https://github.com/jonniedie/ComponentArrays.jl) to wrap several objects inside a single `ComponentVector`.
+
+### Batched evaluation
+
+This is not supported at the moment, but we plan to allow several pushforward / pullback seeds at once (similar to the chunking in ForwardDiff.jl or the batches in Enzyme.jl).
diff --git a/DifferentiationInterface/docs/src/tutorial.md b/DifferentiationInterface/docs/src/tutorial.md
@@ -8,7 +8,7 @@ We present a typical workflow with DifferentiationInterface.jl and showcase its
 
 ```@repl tuto
 using DifferentiationInterface
-import ForwardDiff, Enzyme
+import ADTypes, ForwardDiff, Enzyme
 using BenchmarkTools
 ```
 
@@ -32,7 +32,7 @@ Most backend types are defined by [ADTypes.jl](https://github.com/SciML/ADTypes.
 [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl) is very generic and efficient for low-dimensional inputs, so it's a good starting point:
 
 ```@repl tuto
-backend = AutoForwardDiff()
+backend = ADTypes.AutoForwardDiff()
 ```
 
 Now you can use DifferentiationInterface.jl to get the gradient:
@@ -116,7 +116,7 @@ So let's try the state-of-the-art [Enzyme.jl](https://github.com/EnzymeAD/Enzyme
 For this one, the backend definition is slightly more involved, because you need to feed the "mode" to the object from ADTypes.jl:
 
 ```@repl tuto
-backend2 = AutoEnzyme(Enzyme.Reverse)
+backend2 = ADTypes.AutoEnzyme(; mode=Enzyme.Reverse)
 ```
 
 But once it is done, things run smoothly with exactly the same syntax:

diff --git a/...xt/DifferentiationInterfaceChainRulesCoreExt/DifferentiationInterfaceChainRulesCoreExt.jl b/...xt/DifferentiationInterfaceChainRulesCoreExt/DifferentiationInterfaceChainRulesCoreExt.jl
@@ -13,8 +13,6 @@ const AutoReverseChainRules = AutoChainRules{<:RuleConfig{>:HasReverseMode}}
 
 DI.check_available(::AutoChainRules) = true
 DI.mutation_support(::AutoChainRules) = DI.MutationNotSupported()
-DI.mode(::AutoForwardChainRules) = ADTypes.AbstractForwardMode
-DI.mode(::AutoReverseChainRules) = ADTypes.AbstractReverseMode
 
 ## Pullback