Use weak dependencies if supported

.github/workflows/CI.yml

@@ -20,16 +20,15 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.0'
           - '1.6'
           - '1'
-          #- 'nightly'
+          - 'nightly'
         os:
           - ubuntu-latest
         arch:
           - x64
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
       - uses: julia-actions/setup-julia@v1
         with:
           version: ${{ matrix.version }}
@@ -38,6 +37,6 @@ jobs:
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
       - uses: julia-actions/julia-processcoverage@v1
-      - uses: codecov/codecov-action@v1
+      - uses: codecov/codecov-action@v3
         with:
-          file: lcov.info
+          files: lcov.info

Project.toml

@@ -1,25 +1,46 @@
 name = "AbstractDifferentiation"
 uuid = "c29ec348-61ec-40c8-8164-b8c60e9d9f3d"
 authors = ["Mohamed Tarek <mohamed82008@gmail.com> and contributors"]
-version = "0.4.4"
+version = "0.5.0"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 ExprTools = "e2ba6199-217a-4e67-a87a-7c52f15ade04"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 
+[weakdeps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
+FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
+Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+
+[extensions]
+AbstractDifferentiationChainRulesCoreExt = "ChainRulesCore"
+AbstractDifferentiationFiniteDifferencesExt = "FiniteDifferences"
+AbstractDifferentiationForwardDiffExt = ["DiffResults", "ForwardDiff"]
+AbstractDifferentiationReverseDiffExt = ["DiffResults", "ReverseDiff"]
+AbstractDifferentiationTrackerExt = "Tracker"
+AbstractDifferentiationZygoteExt = "Zygote"
+
 [compat]
 ChainRulesCore = "1"
-Compat = "3, 4"
+DiffResults = "1"
 ExprTools = "0.1"
+FiniteDifferences = "0.12"
 ForwardDiff = "0.10"
-Requires = "0.5, 1"
+Requires = "1"
 ReverseDiff = "1"
-julia = "1"
+Tracker = "0.2"
+Zygote = "0.6"
+julia = "1.6"
 
 [extras]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -29,4 +50,4 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Test", "FiniteDifferences", "ForwardDiff", "Random", "ReverseDiff", "Tracker", "Zygote"]
+test = ["Test", "ChainRulesCore", "DiffResults", "FiniteDifferences", "ForwardDiff", "Random", "ReverseDiff", "Tracker", "Zygote"]

README.md

@@ -1,6 +1,6 @@
 # AbstractDifferentiation
 
-[![CI](https://github.com/JuliaDiff/AbstractDifferentiation.jl/workflows/CI/badge.svg?branch=master)](https://github.com/JuliaDiff/AbstractDifferentiation.jl/actions?query=workflow%3ACI)
+[![CI](https://github.com/JuliaDiff/AbstractDifferentiation.jl/actions/workflows/CI.yml/badge.svg?branch=master)](https://github.com/JuliaDiff/AbstractDifferentiation.jl/actions/workflows/CI.yml?query=branch%3Amaster)
 [![Coverage](https://codecov.io/gh/JuliaDiff/AbstractDifferentiation.jl/branch/master/graph/badge.svg)](https://codecov.io/gh/JuliaDiff/AbstractDifferentiation.jl)
 
 ## Motivation
@@ -11,18 +11,18 @@ Julia has more (automatic) differentiation packages than you can count on 2 hand
 
 ## Loading `AbstractDifferentiation`
 
-To load `AbstractDifferentiation`, use:
+To load `AbstractDifferentiation`, it is recommended to use
 ```julia
-using AbstractDifferentiation
+import AbstractDifferentiation as AD
 ```
-`AbstractDifferentiation` exports a single name `AD` which is just an alias for the `AbstractDifferentiation` module itself. You can use this to access names inside `AbstractDifferentiation` using `AD.<>` instead of typing the long name `AbstractDifferentiation`.
+With the `AD` alias you can access names inside of `AbstractDifferentiation` using `AD.<>` instead of typing the long name `AbstractDifferentiation`.
 
 ## `AbstractDifferentiation` backends
 
 To use `AbstractDifferentiation`, first construct a backend instance `ab::AD.AbstractBackend` using your favorite differentiation package in Julia that supports `AbstractDifferentiation`.
 In particular, you may want to use `AD.ReverseRuleConfigBackend(ruleconfig)` for any [ChainRules.jl](https://github.com/JuliaDiff/ChainRules.jl)-compatible reverse mode differentiation package.
 
-The following backends are temporarily made available by `AbstractDifferentiation` as soon as their corresponding package is loaded (thanks to [Requires.jl](https://github.com/JuliaPackaging/Requires.jl)):
+The following backends are temporarily made available by `AbstractDifferentiation` as soon as their corresponding package is loaded (thanks to [weak dependencies](https://pkgdocs.julialang.org/dev/creating-packages/#Weak-dependencies) on Julia ≥ 1.9 and [Requires.jl](https://github.com/JuliaPackaging/Requires.jl) on older Julia versions):
 
 - `AD.ForwardDiffBackend()` for [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl)
 - `AD.FiniteDifferencesBackend()` for [FiniteDifferences.jl](https://github.com/JuliaDiff/FiniteDifferences.jl)
@@ -35,7 +35,7 @@ In the long term, these backend objects (and many more) will be defined within t
 Here's an example:
 
 ```julia
-julia> using AbstractDifferentiation, Zygote
+julia> import AbstractDifferentiation as AD, Zygote
 
 julia> ab = AD.ZygoteBackend()
 AbstractDifferentiation.ReverseRuleConfigBackend{Zygote.ZygoteRuleConfig{Zygote.Context}}(Zygote.ZygoteRuleConfig{Zygote.Context}(Zygote.Context(nothing)))

ext/AbstractDifferentiationChainRulesCoreExt.jl

@@ -0,0 +1,13 @@
+module AbstractDifferentiationChainRulesCoreExt
+
+import AbstractDifferentiation as AD
+using ChainRulesCore: ChainRulesCore
+
+AD.@primitive function pullback_function(ba::AD.ReverseRuleConfigBackend, f, xs...)
+    _, back = ChainRulesCore.rrule_via_ad(ba.ruleconfig, f, xs...)
+    pullback(vs) = Base.tail(back(vs))
+    pullback(vs::Tuple{Any}) = Base.tail(back(first(vs)))
+    return pullback
+end
+
+end # module

ext/AbstractDifferentiationFiniteDifferencesExt.jl

@@ -0,0 +1,34 @@
+module AbstractDifferentiationFiniteDifferencesExt
+
+import AbstractDifferentiation as AD
+if AD.EXTENSIONS_SUPPORTED
+    using FiniteDifferences: FiniteDifferences
+else
+    using ..FiniteDifferences: FiniteDifferences
+end
+
+"""
+    FiniteDifferencesBackend(method=FiniteDifferences.central_fdm(5, 1))
+
+Create an AD backend that uses forward mode with FiniteDifferences.jl.
+"""
+AD.FiniteDifferencesBackend() = AD.FiniteDifferencesBackend(FiniteDifferences.central_fdm(5, 1))
+
+AD.@primitive function jacobian(ba::AD.FiniteDifferencesBackend, f, xs...)
+    return FiniteDifferences.jacobian(ba.method, f, xs...)
+end
+
+function AD.pushforward_function(ba::AD.FiniteDifferencesBackend, f, xs...)
+    return function pushforward(vs)
+        ws = FiniteDifferences.jvp(ba.method, f, tuple.(xs, vs)...)
+        return length(xs) == 1 ? (ws,) : ws
+    end
+end
+
+function AD.pullback_function(ba::AD.FiniteDifferencesBackend, f, xs...)
+    function pullback(vs)
+        return FiniteDifferences.j′vp(ba.method, f, vs, xs...)
+    end
+end
+
+end # module

src/forwarddiff.jl → ext/AbstractDifferentiationForwardDiffExt.jl

@@ -1,16 +1,13 @@
-using .ForwardDiff: ForwardDiff, DiffResults
-
-"""
-    ForwardDiffBackend{CS}
-
-AD backend that uses forward mode with ForwardDiff.jl.
-
-The type parameter `CS` denotes the chunk size of the differentiation algorithm. If it is
-`Nothing`, then ForwardiffDiff uses a heuristic to set the chunk size based on the input.
-
-See also: [ForwardDiff.jl: Configuring Chunk Size](https://juliadiff.org/ForwardDiff.jl/dev/user/advanced/#Configuring-Chunk-Size)
-"""
-struct ForwardDiffBackend{CS} <: AbstractForwardMode end
+module AbstractDifferentiationForwardDiffExt
+
+import AbstractDifferentiation as AD
+if AD.EXTENSIONS_SUPPORTED
+    using DiffResults: DiffResults
+    using ForwardDiff: ForwardDiff
+else
+    using ..DiffResults: DiffResults
+    using ..ForwardDiff: ForwardDiff
+end
 
 """
     ForwardDiffBackend(; chunksize::Union{Val,Nothing}=nothing)
@@ -23,11 +20,11 @@ ForwarddDiff uses a heuristic to set the chunk size based on the input. Alternat
 
 See also: [ForwardDiff.jl: Configuring Chunk Size](https://juliadiff.org/ForwardDiff.jl/dev/user/advanced/#Configuring-Chunk-Size)
 """
-function ForwardDiffBackend(; chunksize::Union{Val,Nothing}=nothing)
-    return ForwardDiffBackend{getchunksize(chunksize)}()
+function AD.ForwardDiffBackend(; chunksize::Union{Val,Nothing}=nothing)
+    return AD.ForwardDiffBackend{getchunksize(chunksize)}()
 end
 
-@primitive function pushforward_function(ba::ForwardDiffBackend, f, xs...)
+AD.@primitive function pushforward_function(ba::AD.ForwardDiffBackend, f, xs...)
     return function pushforward(vs)
         if length(xs) == 1
             v = vs isa Tuple ? only(vs) : vs
@@ -38,35 +35,35 @@ end
     end
 end
 
-primal_value(x::ForwardDiff.Dual) = ForwardDiff.value(x)
-primal_value(x::AbstractArray{<:ForwardDiff.Dual}) = ForwardDiff.value.(x)
+AD.primal_value(x::ForwardDiff.Dual) = ForwardDiff.value(x)
+AD.primal_value(x::AbstractArray{<:ForwardDiff.Dual}) = ForwardDiff.value.(x)
 
 # these implementations are more efficient than the fallbacks
 
-function gradient(ba::ForwardDiffBackend, f, x::AbstractArray)
+function AD.gradient(ba::AD.ForwardDiffBackend, f, x::AbstractArray)
     cfg = ForwardDiff.GradientConfig(f, x, chunk(ba, x))
     return (ForwardDiff.gradient(f, x, cfg),)
 end
 
-function jacobian(ba::ForwardDiffBackend, f, x::AbstractArray)
-    cfg = ForwardDiff.JacobianConfig(asarray ∘ f, x, chunk(ba, x))
-    return (ForwardDiff.jacobian(asarray ∘ f, x, cfg),)
+function AD.jacobian(ba::AD.ForwardDiffBackend, f, x::AbstractArray)
+    cfg = ForwardDiff.JacobianConfig(AD.asarray ∘ f, x, chunk(ba, x))
+    return (ForwardDiff.jacobian(AD.asarray ∘ f, x, cfg),)
 end
-jacobian(::ForwardDiffBackend, f, x::Number) = (ForwardDiff.derivative(f, x),)
+AD.jacobian(::AD.ForwardDiffBackend, f, x::Number) = (ForwardDiff.derivative(f, x),)
 
-function hessian(ba::ForwardDiffBackend, f, x::AbstractArray)
+function AD.hessian(ba::AD.ForwardDiffBackend, f, x::AbstractArray)
     cfg = ForwardDiff.HessianConfig(f, x, chunk(ba, x))
     return (ForwardDiff.hessian(f, x, cfg),)
 end
 
-function value_and_gradient(ba::ForwardDiffBackend, f, x::AbstractArray)
+function AD.value_and_gradient(ba::AD.ForwardDiffBackend, f, x::AbstractArray)
     result = DiffResults.GradientResult(x)
     cfg = ForwardDiff.GradientConfig(f, x, chunk(ba, x))
     ForwardDiff.gradient!(result, f, x, cfg)
     return DiffResults.value(result), (DiffResults.derivative(result),)
 end
 
-function value_and_hessian(ba::ForwardDiffBackend, f, x)
+function AD.value_and_hessian(ba::AD.ForwardDiffBackend, f, x)
     result = DiffResults.HessianResult(x)
     cfg = ForwardDiff.HessianConfig(f, result, x, chunk(ba, x))
     ForwardDiff.hessian!(result, f, x, cfg)
@@ -80,5 +77,7 @@ end
 getchunksize(::Nothing) = Nothing
 getchunksize(::Val{N}) where {N} = N
 
-chunk(::ForwardDiffBackend{Nothing}, x) = ForwardDiff.Chunk(x)
-chunk(::ForwardDiffBackend{N}, _) where {N} = ForwardDiff.Chunk{N}()
+chunk(::AD.ForwardDiffBackend{Nothing}, x) = ForwardDiff.Chunk(x)
+chunk(::AD.ForwardDiffBackend{N}, _) where {N} = ForwardDiff.Chunk{N}()
+
+end # module

ext/AbstractDifferentiationReverseDiffExt.jl

@@ -0,0 +1,64 @@
+module AbstractDifferentiationReverseDiffExt
+
+import AbstractDifferentiation as AD
+if AD.EXTENSIONS_SUPPORTED
+    using DiffResults: DiffResults
+    using ReverseDiff: ReverseDiff
+else
+    using ..DiffResults: DiffResults
+    using ..ReverseDiff: ReverseDiff
+end
+
+AD.primal_value(x::ReverseDiff.TrackedReal) = ReverseDiff.value(x)
+AD.primal_value(x::AbstractArray{<:ReverseDiff.TrackedReal}) = ReverseDiff.value.(x)
+AD.primal_value(x::ReverseDiff.TrackedArray) = ReverseDiff.value(x)
+
+AD.@primitive function jacobian(ba::AD.ReverseDiffBackend, f, xs...)
+    xs_arr = map(AD.asarray, xs)
+    tape = ReverseDiff.JacobianTape(xs_arr) do (xs_arr...)
+        xs_new = map(xs, xs_arr) do x, x_arr
+            return x isa Number ? only(x_arr) : x_arr
+        end
+        return AD.asarray(f(xs_new...))
+    end
+    results = ReverseDiff.jacobian!(tape, xs_arr)
+    return map(xs, results) do x, result
+        return x isa Number ? vec(result) : result
+    end
+end
+function AD.jacobian(::AD.ReverseDiffBackend, f, xs::AbstractArray...)
+    return ReverseDiff.jacobian(AD.asarray ∘ f, xs)
+end
+
+function AD.derivative(::AD.ReverseDiffBackend, f, xs::Number...)
+    tape = ReverseDiff.InstructionTape()
+    xs_tracked = ReverseDiff.TrackedReal.(xs, zero.(xs), Ref(tape))
+    y_tracked = f(xs_tracked...)
+    ReverseDiff.seed!(y_tracked)
+    ReverseDiff.reverse_pass!(tape)
+    return ReverseDiff.deriv.(xs_tracked)
+end
+
+function AD.gradient(::AD.ReverseDiffBackend, f, xs::AbstractArray...)
+    return ReverseDiff.gradient(f, xs)
+end
+
+function AD.hessian(::AD.ReverseDiffBackend, f, x::AbstractArray)
+    return (ReverseDiff.hessian(f, x),)
+end
+
+function AD.value_and_gradient(::AD.ReverseDiffBackend, f, x::AbstractArray)
+    result = DiffResults.GradientResult(x)
+    cfg = ReverseDiff.GradientConfig(x)
+    ReverseDiff.gradient!(result, f, x, cfg)
+    return DiffResults.value(result), (DiffResults.derivative(result),)
+end
+
+function AD.value_and_hessian(::AD.ReverseDiffBackend, f, x)
+    result = DiffResults.HessianResult(x)
+    cfg = ReverseDiff.HessianConfig(result, x)
+    ReverseDiff.hessian!(result, f, x, cfg)
+    return DiffResults.value(result), (DiffResults.hessian(result),)
+end
+
+end # module

ext/AbstractDifferentiationTrackerExt.jl

@@ -0,0 +1,38 @@
+module AbstractDifferentiationTrackerExt
+
+import AbstractDifferentiation as AD
+if AD.EXTENSIONS_SUPPORTED
+    using Tracker: Tracker
+else
+    using ..Tracker: Tracker
+end
+
+function AD.second_lowest(::AD.TrackerBackend)
+    return throw(ArgumentError("Tracker backend does not support nested differentiation."))
+end
+
+AD.primal_value(x::Tracker.TrackedReal) = Tracker.data(x)
+AD.primal_value(x::Tracker.TrackedArray) = Tracker.data(x)
+AD.primal_value(x::AbstractArray{<:Tracker.TrackedReal}) = Tracker.data.(x)
+
+AD.@primitive function pullback_function(ba::AD.TrackerBackend, f, xs...)
+    value, back = Tracker.forward(f, xs...)
+    function pullback(ws)
+        if ws isa Tuple && !(value isa Tuple) 
+            map(Tracker.data, back(only(ws)))
+        else
+            map(Tracker.data, back(ws))
+        end
+    end
+    return pullback
+end
+
+function AD.derivative(::AD.TrackerBackend, f, xs::Number...)
+    return Tracker.data.(Tracker.gradient(f, xs...))
+end
+
+function AD.gradient(::AD.TrackerBackend, f, xs::AbstractVector...)
+    return Tracker.data.(Tracker.gradient(f, xs...))
+end
+
+end # module

ext/AbstractDifferentiationZygoteExt.jl

@@ -0,0 +1,12 @@
+module AbstractDifferentiationZygoteExt
+
+import AbstractDifferentiation as AD
+if AD.EXTENSIONS_SUPPORTED
+    using Zygote: Zygote
+else
+    using ..Zygote: Zygote
+end
+
+AD.ZygoteBackend() = AD.ReverseRuleConfigBackend(Zygote.ZygoteRuleConfig())
+
+end # module