Tutorial
We present a typical workflow with DifferentiationInterfaceTest.jl, building on the tutorial of the DifferentiationInterface.jl documentation (which we encourage you to read first).
julia> using DifferentiationInterface, DifferentiationInterfaceTest
julia> using ForwardDiff: ForwardDiff
julia> using Zygote: Zygote
Introduction
The AD backends we want to compare are ForwardDiff.jl and Zygote.jl.
backends = [AutoForwardDiff(), AutoZygote()]
2-element Vector{ADTypes.AbstractADType}:
AutoForwardDiff()
AutoZygote()
To do that, we are going to take gradients of a simple function:
f(x::AbstractArray) = sum(sin, x)
f (generic function with 1 method)
Of course we know the true gradient mapping:
∇f(x::AbstractArray) = cos.(x)
∇f (generic function with 1 method)
DifferentiationInterfaceTest.jl relies with so-called Scenario
s, in which you encapsulate the information needed for your test:
- the operator category (here
:gradient
) - the behavior of the operator (either
:in
or:out
of place) - the function
f
- the input
x
of the functionf
(and possible tangents or contexts) - the reference first-order result
res1
(and possible second-order resultres2
) of the operator - the arguments
prep_args
passed during preparation
xv = rand(Float32, 3)
xm = rand(Float64, 3, 2)
scenarios = [
Scenario{:gradient,:out}(f, xv; res1=∇f(xv)),
Scenario{:gradient,:out}(f, xm; res1=∇f(xm)),
];
Testing
The main entry point for testing is the function test_differentiation
. It has many options, but the main ingredients are the following:
julia> test_differentiation( backends, # the backends you want to compare scenarios; # the scenarios you defined, correctness=true, # compares values against the reference type_stability=:none, # checks type stability with JET.jl detailed=true, # prints a detailed test set )
Test Summary: | Pass Total Time Testing correctness | 88 88 11.3s AutoForwardDiff() | 44 44 5.4s gradient | 44 44 5.3s Scenario{:gradient,:out} f : Vector{Float32} -> Float32 | 22 22 2.5s Scenario{:gradient,:out} f : Matrix{Float64} -> Float64 | 22 22 1.6s AutoZygote() | 44 44 5.8s gradient | 44 44 5.8s Scenario{:gradient,:out} f : Vector{Float32} -> Float32 | 22 22 4.7s Scenario{:gradient,:out} f : Matrix{Float64} -> Float64 | 22 22 1.1s
Benchmarking
Once you are confident that your backends give the correct answers, you probably want to compare their performance. This is made easy by the benchmark_differentiation
function, whose syntax should feel familiar:
df = benchmark_differentiation(backends, scenarios);
Row | backend | scenario | operator | prepared | calls | samples | evals | time | allocs | bytes | gc_fraction | compile_fraction |
---|---|---|---|---|---|---|---|---|---|---|---|---|
Abstract… | Scenario… | Symbol | Bool | Int64 | Int64 | Int64 | Float64 | Float64 | Float64 | Float64 | Float64 | |
1 | AutoForwardDiff() | Scenario{:gradient,:out} f : Vector{Float32} -> Float32 | value_and_gradient | true | 1 | 25846 | 533 | 5.28555e-8 | 3.0 | 112.0 | 0.0 | 0.0 |
2 | AutoForwardDiff() | Scenario{:gradient,:out} f : Vector{Float32} -> Float32 | gradient | true | 1 | 28128 | 614 | 4.48404e-8 | 2.0 | 80.0 | 0.0 | 0.0 |
3 | AutoForwardDiff() | Scenario{:gradient,:out} f : Matrix{Float64} -> Float64 | value_and_gradient | true | 1 | 40273 | 156 | 1.21506e-7 | 3.0 | 160.0 | 0.0 | 0.0 |
4 | AutoForwardDiff() | Scenario{:gradient,:out} f : Matrix{Float64} -> Float64 | gradient | true | 1 | 33443 | 212 | 1.11953e-7 | 2.0 | 128.0 | 0.0 | 0.0 |
5 | AutoZygote() | Scenario{:gradient,:out} f : Vector{Float32} -> Float32 | value_and_gradient | true | 1 | 28231 | 34 | 8.28e-7 | 25.0 | 688.0 | 0.0 | 0.0 |
6 | AutoZygote() | Scenario{:gradient,:out} f : Vector{Float32} -> Float32 | gradient | true | 1 | 31055 | 40 | 6.24925e-7 | 23.0 | 624.0 | 0.0 | 0.0 |
7 | AutoZygote() | Scenario{:gradient,:out} f : Matrix{Float64} -> Float64 | value_and_gradient | true | 1 | 122581 | 3 | 1.11867e-6 | 29.0 | 1040.0 | 0.0 | 0.0 |
8 | AutoZygote() | Scenario{:gradient,:out} f : Matrix{Float64} -> Float64 | gradient | true | 1 | 27915 | 30 | 9.117e-7 | 27.0 | 976.0 | 0.0 | 0.0 |
The resulting object is a DataFrame
from DataFrames.jl, whose columns correspond to the fields of DifferentiationBenchmarkDataRow
.