Tutorial
We present a typical workflow with DifferentiationInterfaceTest.jl, building on the DifferentiationInterface.jl tutorial (which we encourage you to read first).
julia> using DifferentiationInterface, DifferentiationInterfaceTest
julia> import ForwardDiff, Enzyme
julia> import DataFrames, Markdown, PrettyTables, Printf
Introduction
The AD backends we want to compare are ForwardDiff.jl and Enzyme.jl.
julia> backends = [AutoForwardDiff(), AutoEnzyme(; mode=Enzyme.Reverse)]
2-element Vector{ADTypes.AbstractADType}: AutoForwardDiff{nothing, Nothing}(nothing) AutoEnzyme{EnzymeCore.ReverseMode{false, EnzymeCore.FFIABI, false}}(EnzymeCore.ReverseMode{false, EnzymeCore.FFIABI, false}())
To do that, we are going to take gradients of a simple function:
julia> f(x::AbstractArray) = sum(sin, x)
f (generic function with 1 method)
Of course we know the true gradient mapping:
julia> ∇f(x::AbstractArray) = cos.(x)
∇f (generic function with 1 method)
DifferentiationInterfaceTest.jl relies with so-called "scenarios", in which you encapsulate the information needed for your test:
- the function
f
- the input
x
(and outputy
for mutating functions) - optionally a reference
ref
to check against
There is one scenario per operator, and so here we will use GradientScenario
:
julia> scenarios = [ GradientScenario(f; x=rand(Float32, 3), ref=∇f, place=:inplace), GradientScenario(f; x=rand(Float64, 3, 2), ref=∇f, place=:inplace) ];
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=false, # checks type stability with JET.jl detailed=true, # prints a detailed test set )
Test Summary: | Pass Total Time Testing correctness | 32 32 36.6s ForwardDiff (forward) | 16 16 3.7s GradientScenario | 16 16 3.7s GradientScenario{1,inplace}(f : Vector{Float32} -> Float32) | 8 8 1.9s GradientScenario{1,inplace}(f : Matrix{Float64} -> Float64) | 8 8 1.8s Enzyme (reverse) | 16 16 32.9s GradientScenario | 16 16 32.9s GradientScenario{1,inplace}(f : Vector{Float32} -> Float32) | 8 8 31.4s GradientScenario{1,inplace}(f : Matrix{Float64} -> Float64) | 8 8 1.4s
If you are too lazy to manually specify the reference, you can also provide an AD backend as the ref_backend
keyword argument, which will serve as the ground truth for comparison.
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:
julia> benchmark_result = benchmark_differentiation(backends, scenarios);
The resulting object is a Vector
of DifferentiationBenchmarkDataRow
, which can easily be converted into a DataFrame
from DataFrames.jl:
julia> df = DataFrames.DataFrame(benchmark_result)
12×11 DataFrame Row │ backend scenario o ⋯ │ Abstract… Gradient… S ⋯ ─────┼────────────────────────────────────────────────────────────────────────── 1 │ AutoForwardDiff{nothing, Nothing… GradientScenario{1,inplace}(f : … p ⋯ 2 │ AutoForwardDiff{nothing, Nothing… GradientScenario{1,inplace}(f : … v 3 │ AutoForwardDiff{nothing, Nothing… GradientScenario{1,inplace}(f : … g 4 │ AutoForwardDiff{nothing, Nothing… GradientScenario{1,inplace}(f : … p 5 │ AutoForwardDiff{nothing, Nothing… GradientScenario{1,inplace}(f : … v ⋯ 6 │ AutoForwardDiff{nothing, Nothing… GradientScenario{1,inplace}(f : … g 7 │ AutoEnzyme{ReverseMode{false, FF… GradientScenario{1,inplace}(f : … p 8 │ AutoEnzyme{ReverseMode{false, FF… GradientScenario{1,inplace}(f : … v 9 │ AutoEnzyme{ReverseMode{false, FF… GradientScenario{1,inplace}(f : … g ⋯ 10 │ AutoEnzyme{ReverseMode{false, FF… GradientScenario{1,inplace}(f : … p 11 │ AutoEnzyme{ReverseMode{false, FF… GradientScenario{1,inplace}(f : … v 12 │ AutoEnzyme{ReverseMode{false, FF… GradientScenario{1,inplace}(f : … g 9 columns omitted
Here's what the resulting DataFrame
looks like with all its columns. Note that we only compare (possibly) in-place operators, because they are always more efficient.
table = PrettyTables.pretty_table(
String,
df;
backend=Val(:markdown),
header=names(df),
)
Markdown.parse(table)
backend | scenario | operator | calls | samples | evals | time | allocs | bytes | gc_fraction | compile_fraction |
---|---|---|---|---|---|---|---|---|---|---|
AutoForwardDiff{nothing, Nothing}(nothing) | GradientScenario{1,inplace}(f : Vector{Float32} -> Float32) | prepare_gradient | 0 | 1 | 1 | 1.1261e-5 | 11.0 | 528.0 | 0.0 | 0.0 |
AutoForwardDiff{nothing, Nothing}(nothing) | GradientScenario{1,inplace}(f : Vector{Float32} -> Float32) | value_and_gradient! | 1 | 34869 | 1 | 6.9e-8 | 1.0 | 32.0 | 0.0 | 0.0 |
AutoForwardDiff{nothing, Nothing}(nothing) | GradientScenario{1,inplace}(f : Vector{Float32} -> Float32) | gradient! | 1 | 42507 | 1 | 6.0e-8 | 0.0 | 0.0 | 0.0 | 0.0 |
AutoForwardDiff{nothing, Nothing}(nothing) | GradientScenario{1,inplace}(f : Matrix{Float64} -> Float64) | prepare_gradient | 0 | 1 | 1 | 1.4417e-5 | 11.0 | 1776.0 | 0.0 | 0.0 |
AutoForwardDiff{nothing, Nothing}(nothing) | GradientScenario{1,inplace}(f : Matrix{Float64} -> Float64) | value_and_gradient! | 1 | 24425 | 1 | 2.0e-7 | 5.0 | 192.0 | 0.0 | 0.0 |
AutoForwardDiff{nothing, Nothing}(nothing) | GradientScenario{1,inplace}(f : Matrix{Float64} -> Float64) | gradient! | 1 | 27782 | 1 | 2.0e-7 | 4.0 | 160.0 | 0.0 | 0.0 |
AutoEnzyme{ReverseMode{false, FFIABI, false}}(ReverseMode{false, FFIABI, false}()) | GradientScenario{1,inplace}(f : Vector{Float32} -> Float32) | prepare_gradient | 0 | 1 | 1 | 2.4e-7 | 0.0 | 0.0 | 0.0 | 0.0 |
AutoEnzyme{ReverseMode{false, FFIABI, false}}(ReverseMode{false, FFIABI, false}()) | GradientScenario{1,inplace}(f : Vector{Float32} -> Float32) | value_and_gradient! | 1 | 52083 | 1 | 8.91e-7 | 9.0 | 192.0 | 0.0 | 0.0 |
AutoEnzyme{ReverseMode{false, FFIABI, false}}(ReverseMode{false, FFIABI, false}()) | GradientScenario{1,inplace}(f : Vector{Float32} -> Float32) | gradient! | 1 | 78381 | 1 | 4.9e-8 | 0.0 | 0.0 | 0.0 | 0.0 |
AutoEnzyme{ReverseMode{false, FFIABI, false}}(ReverseMode{false, FFIABI, false}()) | GradientScenario{1,inplace}(f : Matrix{Float64} -> Float64) | prepare_gradient | 0 | 1 | 1 | 2.2e-7 | 0.0 | 0.0 | 0.0 | 0.0 |
AutoEnzyme{ReverseMode{false, FFIABI, false}}(ReverseMode{false, FFIABI, false}()) | GradientScenario{1,inplace}(f : Matrix{Float64} -> Float64) | value_and_gradient! | 1 | 24254 | 1 | 1.252e-6 | 9.0 | 192.0 | 0.0 | 0.0 |
AutoEnzyme{ReverseMode{false, FFIABI, false}}(ReverseMode{false, FFIABI, false}()) | GradientScenario{1,inplace}(f : Matrix{Float64} -> Float64) | gradient! | 1 | 52911 | 1 | 6.9e-8 | 0.0 | 0.0 | 0.0 | 0.0 |