Backend-specific utilities (derivative, multiderivative, gradient, jacobian)

.gitignore

@@ -6,6 +6,7 @@
 /benchmark/*.json
 
 /docs/build/
+/docs/src/index.md
 
 /Manifest.toml
 /docs/Manifest.toml

Project.toml

@@ -14,6 +14,7 @@ DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
@@ -22,6 +23,7 @@ DifferentiationInterfaceChainRulesCoreExt = "ChainRulesCore"
 DifferentiationInterfaceEnzymeExt = "Enzyme"
 DifferentiationInterfaceFiniteDiffExt = "FiniteDiff"
 DifferentiationInterfaceForwardDiffExt = ["ForwardDiff", "DiffResults"]
+DifferentiationInterfacePolyesterForwardDiffExt = ["PolyesterForwardDiff", "ForwardDiff", "DiffResults"]
 DifferentiationInterfaceReverseDiffExt = ["ReverseDiff", "DiffResults"]
 DifferentiationInterfaceZygoteExt = ["Zygote"]
 
@@ -34,6 +36,7 @@ FillArrays = "1"
 FiniteDiff = "2.22"
 ForwardDiff = "0.10"
 LinearAlgebra = "1"
+PolyesterForwardDiff = "0.1"
 ReverseDiff = "1.15"
 Zygote = "0.6"
 julia = "1.10"

README.md

@@ -5,6 +5,57 @@
 [![Coverage](https://codecov.io/gh/gdalle/DifferentiationInterface.jl/branch/main/graph/badge.svg)](https://codecov.io/gh/gdalle/DifferentiationInterface.jl)
 [![Code Style: Blue](https://img.shields.io/badge/code%20style-blue-4495d1.svg)](https://github.com/invenia/BlueStyle)
 
-An experimental redesign for [AbstractDifferentiation.jl](https://github.com/JuliaDiff/AbstractDifferentiation.jl).
+An interface to various automatic differentiation backends in Julia.
 
-See the documentation for details.
+## Goal
+
+This package provides a backend-agnostic syntax to differentiate functions `f(x) = y`, where `x` and `y` are either numbers or abstract arrays.
+
+It started out as an experimental redesign for [AbstractDifferentiation.jl](https://github.com/JuliaDiff/AbstractDifferentiation.jl).
+
+## Example
+
+```jldoctest
+julia> using DifferentiationInterface, Enzyme
+
+julia> backend = EnzymeReverseBackend();
+
+julia> f(x) = sum(abs2, x);
+
+julia> value_and_gradient(backend, f, [1., 2., 3.])
+(14.0, [2.0, 4.0, 6.0])
+```
+
+## Design
+
+Each backend must implement only one primitive:
+
+- forward mode: the pushforward, computing a Jacobian-vector product
+- reverse mode: the pullback, computing a vector-Jacobian product
+
+From these primitives, several utilities are defined, depending on the type of the input and output:
+
+|              | scalar output | array output    |
+| ------------ | ------------- | --------------- |
+| scalar input | derivative    | multiderivative |
+| array input  | gradient      | jacobian        |
+
+## Supported backends
+
+Forward mode:
+
+- [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl)
+- [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl)
+- [ChainRulesCore.jl](https://github.com/JuliaDiff/ChainRulesCore.jl)
+- [FiniteDiff.jl](https://github.com/JuliaDiff/FiniteDiff.jl)
+
+Reverse mode:
+
+- [ChainRulesCore.jl](https://github.com/JuliaDiff/ChainRulesCore.jl)
+- [Zygote.jl](https://github.com/FluxML/Zygote.jl)
+- [ReverseDiff.jl](https://github.com/JuliaDiff/ReverseDiff.jl)
+
+Experimental:
+
+- [PolyesterForwardDiff.jl](https://github.com/JuliaDiff/PolyesterForwardDiff.jl)
+- [Diffractor.jl](https://github.com/JuliaDiff/Diffractor.jl) (currently broken due to [#277](https://github.com/JuliaDiff/Diffractor.jl/issues/277))

benchmark/Project.toml

@@ -6,5 +6,6 @@ Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

benchmark/benchmarks.jl

@@ -1,63 +1,210 @@
-# Run benchmarks locally by calling:
-# julia -e 'using BenchmarkCI; BenchmarkCI.judge(baseline="origin/main"); BenchmarkCI.displayjudgement()'
-
-using Base: Fix2
 using BenchmarkTools
 using DifferentiationInterface
 using LinearAlgebra
 
-using Diffractor: Diffractor
 using Enzyme: Enzyme
 using FiniteDiff: FiniteDiff
 using ForwardDiff: ForwardDiff
+using PolyesterForwardDiff: PolyesterForwardDiff
 using ReverseDiff: ReverseDiff
 using Zygote: Zygote
 
-scalar_to_scalar(x::Real) = x
-scalar_to_vector(x::Real, n) = collect((1:n) .* x)
-vector_to_scalar(x::AbstractVector{<:Real}) = dot(1:length(x), x)
-vector_to_vector(x::AbstractVector{<:Real}) = (1:length(x)) .* x
+## Settings
+
+BenchmarkTools.DEFAULT_PARAMETERS.evals = 1
+BenchmarkTools.DEFAULT_PARAMETERS.samples = 100
+BenchmarkTools.DEFAULT_PARAMETERS.seconds = 1
+
+## Functions
+
+struct Layer{W<:Union{Number,AbstractArray},B<:Union{Number,AbstractArray},S<:Function}
+    w::W
+    b::B
+    σ::S
+end
+
+function (l::Layer{<:Number,<:Number})(x::Number)::Number
+    return l.σ(l.w * x + l.b)
+end
+
+function (l::Layer{<:AbstractVector,<:AbstractVector})(x::Number)::AbstractVector
+    return l.σ.(l.w .* x .+ l.b)
+end
 
-forward_backends = [EnzymeForwardBackend(), FiniteDiffBackend(), ForwardDiffBackend()]
+function (l::Layer{<:AbstractVector,<:Number})(x::AbstractVector)::Number
+    return l.σ(dot(l.w, x) + l.b)
+end
+
+function (l::Layer{<:AbstractMatrix,<:AbstractVector})(x::AbstractVector)::AbstractVector
+    return l.σ.(l.w * x .+ l.b)
+end
 
-reverse_backends = [
-    ChainRulesReverseBackend(Zygote.ZygoteRuleConfig()),
-    EnzymeReverseBackend(),
-    ReverseDiffBackend(),
+## Backends
+
+forward_custom_backends = [
+    EnzymeForwardBackend(; custom=true),
+    FiniteDiffBackend(; custom=true),
+    ForwardDiffBackend(; custom=true),
+    PolyesterForwardDiffBackend(4; custom=true),
+]
+
+forward_fallback_backends = [
+    EnzymeForwardBackend(; custom=false),
+    FiniteDiffBackend(; custom=false),
+    ForwardDiffBackend(; custom=false),
 ]
 
-n_values = [10]
+reverse_custom_backends = [
+    ZygoteBackend(; custom=true),
+    EnzymeReverseBackend(; custom=true),
+    ReverseDiffBackend(; custom=true),
+]
+
+reverse_fallback_backends = [
+    ZygoteBackend(; custom=false),
+    EnzymeReverseBackend(; custom=false),
+    ReverseDiffBackend(; custom=false),
+]
+
+all_custom_backends = vcat(forward_custom_backends, reverse_custom_backends)
+all_fallback_backends = vcat(forward_fallback_backends, reverse_fallback_backends)
+all_backends = vcat(all_custom_backends, all_fallback_backends)
+
+## Suite
 
 SUITE = BenchmarkGroup()
 
-for n in n_values
-    for backend in forward_backends
-        SUITE["forward"]["scalar_to_scalar"][n][string(backend)] = @benchmarkable begin
-            value_and_pushforward!(dy, $backend, scalar_to_scalar, x, dx)
-        end setup = (x = 1.0; dx = 1.0; dy = 0.0) evals = 1
-        if backend != EnzymeForwardBackend()  # type instability?
-            SUITE["forward"]["scalar_to_vector"][n][string(backend)] = @benchmarkable begin
-                value_and_pushforward!(dy, $backend, Fix2(scalar_to_vector, $n), x, dx)
-            end setup = (x = 1.0; dx = 1.0; dy = zeros($n)) evals = 1
+### Scalar to scalar
+
+scalar_to_scalar = Layer(randn(), randn(), tanh)
+
+for backend in all_backends
+    handles_types(backend, Number, Number) || continue
+    SUITE["value_and_derivative"][(1, 1)][string(backend)] = @benchmarkable begin
+        value_and_derivative($backend, $scalar_to_scalar, x)
+    end setup = (x = randn())
+end
+
+for backend in all_fallback_backends
+    handles_types(backend, Number, Number) || continue
+    if autodiff_mode(backend) == :forward
+        SUITE["value_and_pushforward"][(1, 1)][string(backend)] = @benchmarkable begin
+            value_and_pushforward($backend, $scalar_to_scalar, x, dx)
+        end setup = (x = randn(); dx = randn())
+    else
+        SUITE["value_and_pullback"][(1, 1)][string(backend)] = @benchmarkable begin
+            value_and_pullback($backend, $scalar_to_scalar, x, dy)
+        end setup = (x = randn(); dy = randn())
+    end
+end
+
+### Scalar to vector
+
+for m in [10]
+    scalar_to_vector = Layer(randn(m), randn(m), tanh)
+
+    for backend in all_backends
+        handles_types(backend, Number, Vector) || continue
+        SUITE["value_and_multiderivative"][(1, m)][string(backend)] = @benchmarkable begin
+            value_and_multiderivative($backend, $scalar_to_vector, x)
+        end setup = (x = randn())
+        SUITE["value_and_multiderivative!"][(1, m)][string(backend)] = @benchmarkable begin
+            value_and_multiderivative!(multider, $backend, $scalar_to_vector, x)
+        end setup = (x = randn(); multider = zeros($m))
+    end
+
+    for backend in all_fallback_backends
+        handles_types(backend, Number, Vector) || continue
+        if autodiff_mode(backend) == :forward
+            SUITE["value_and_pushforward"][(1, m)][string(backend)] = @benchmarkable begin
+                value_and_pushforward($backend, $scalar_to_vector, x, dx)
+            end setup = (x = randn(); dx = randn())
+            SUITE["value_and_pushforward!"][(1, m)][string(backend)] = @benchmarkable begin
+                value_and_pushforward!(dy, $backend, $scalar_to_vector, x, dx)
+            end setup = (x = randn(); dx = randn(); dy = zeros($m))
+        else
+            SUITE["value_and_pullback"][(1, m)][string(backend)] = @benchmarkable begin
+                value_and_pullback($backend, $scalar_to_vector, x, dy)
+            end setup = (x = randn(); dy = ones($m))
+            SUITE["value_and_pullback!"][(1, m)][string(backend)] = @benchmarkable begin
+                value_and_pullback!(dx, $backend, $scalar_to_vector, x, dy)
+            end setup = (x = randn(); dy = ones($m); dx = 0.0)
         end
-        SUITE["forward"]["vector_to_vector"][n][string(backend)] = @benchmarkable begin
-            value_and_pushforward!(dy, $backend, vector_to_vector, x, dx)
-        end setup = (x = randn($n); dx = randn($n); dy = zeros($n)) evals = 1
+    end
+end
+
+### Vector to scalar
+
+for n in [10]
+    vector_to_scalar = Layer(randn(n), randn(), tanh)
+
+    for backend in all_backends
+        handles_types(backend, Vector, Number) || continue
+        SUITE["value_and_gradient"][(n, 1)][string(backend)] = @benchmarkable begin
+            value_and_gradient($backend, $vector_to_scalar, x)
+        end setup = (x = randn($n))
+        SUITE["value_and_gradient!"][(n, 1)][string(backend)] = @benchmarkable begin
+            value_and_gradient!(grad, $backend, $vector_to_scalar, x)
+        end setup = (x = randn($n); grad = zeros($n))
     end
 
-    for backend in reverse_backends
-        if backend != ReverseDiffBackend()
-            SUITE["reverse"]["scalar_to_scalar"][n][string(backend)] = @benchmarkable begin
-                value_and_pullback!(dx, $backend, scalar_to_scalar, x, dy)
-            end setup = (x = 1.0; dy = 1.0; dx = 0.0) evals = 1
+    for backend in all_fallback_backends
+        handles_types(backend, Vector, Number) || continue
+        if autodiff_mode(backend) == :forward
+            SUITE["value_and_pushforward"][(n, 1)][string(backend)] = @benchmarkable begin
+                value_and_pushforward($backend, $vector_to_scalar, x, dx)
+            end setup = (x = randn($n); dx = randn($n))
+            SUITE["value_and_pushforward!"][(n, 1)][string(backend)] = @benchmarkable begin
+                value_and_pushforward!(dy, $backend, $vector_to_scalar, x, dx)
+            end setup = (x = randn($n); dx = randn($n); dy = 0.0)
+        else
+            SUITE["value_and_pullback"][(n, 1)][string(backend)] = @benchmarkable begin
+                value_and_pullback($backend, $vector_to_scalar, x, dy)
+            end setup = (x = randn($n); dy = randn())
+            SUITE["value_and_pullback!"][(n, 1)][string(backend)] = @benchmarkable begin
+                value_and_pullback!(dx, $backend, $vector_to_scalar, x, dy)
+            end setup = (x = randn($n); dy = randn(); dx = zeros($n))
         end
-        SUITE["reverse"]["vector_to_scalar"][n][string(backend)] = @benchmarkable begin
-            value_and_pullback!(dx, $backend, vector_to_scalar, x, dy)
-        end setup = (x = randn($n); dy = 1.0; dx = zeros($n)) evals = 1
-        if backend != EnzymeReverseBackend()
-            SUITE["reverse"]["vector_to_vector"][n][string(backend)] = @benchmarkable begin
-                value_and_pullback!(dx, $backend, vector_to_vector, x, dy)
-            end setup = (x = randn($n); dy = randn($n); dx = zeros($n)) evals = 1
+    end
+end
+
+### Vector to vector
+
+for (n, m) in [(10, 10)]
+    vector_to_vector = Layer(randn(m, n), randn(m), tanh)
+
+    for backend in all_backends
+        handles_types(backend, Vector, Vector) || continue
+        SUITE["value_and_jacobian"][(n, m)][string(backend)] = @benchmarkable begin
+            value_and_jacobian($backend, $vector_to_vector, x)
+        end setup = (x = randn($n))
+        SUITE["value_and_jacobian!"][(n, m)][string(backend)] = @benchmarkable begin
+            value_and_jacobian!(jac, $backend, $vector_to_vector, x)
+        end setup = (x = randn($n); jac = zeros($m, $n))
+    end
+
+    for backend in all_fallback_backends
+        handles_types(backend, Vector, Vector) || continue
+        if autodiff_mode(backend) == :forward
+            SUITE["value_and_pushforward"][(n, m)][string(backend)] = @benchmarkable begin
+                value_and_pushforward($backend, $vector_to_vector, x, dx)
+            end setup = (x = randn($n); dx = randn($n))
+            SUITE["value_and_pushforward!"][(n, m)][string(backend)] = @benchmarkable begin
+                value_and_pushforward!(dy, $backend, $vector_to_vector, x, dx)
+            end setup = (x = randn($n); dx = randn($n); dy = zeros($m))
+        else
+            SUITE["value_and_pullback"][(n, m)][string(backend)] = @benchmarkable begin
+                value_and_pullback($backend, $vector_to_vector, x, dy)
+            end setup = (x = randn($n); dy = randn($m))
+            SUITE["value_and_pullback!"][(n, m)][string(backend)] = @benchmarkable begin
+                value_and_pullback!(dx, $backend, $vector_to_vector, x, dy)
+            end setup = (x = randn($n); dy = randn($m); dx = zeros($n))
         end
     end
 end
+
+# Run benchmarks locally
+# results = BenchmarkTools.run(SUITE; verbose=true)
+
+# Compare commits locally
+# using BenchmarkCI; BenchmarkCI.judge(baseline="origin/main"); BenchmarkCI.displayjudgement()

docs/Project.toml

@@ -1,11 +1,12 @@
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
-Diffractor = "9f5e2b26-1114-432f-b630-d3fe2085c51c"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"