Revamp test scenarios (#323)

* Revamp test scenarios

* Fix

* Fixes

* Fixes

* Fix

* Refix

* Refix

* Fixes

* Fix

* Fix

* Coverage

* Appease JET

Author: gdalle

.github/workflows/Test.yml

@@ -123,6 +123,7 @@ jobs:
           - Formalities
           - Zero
           - ForwardDiff
+          - Zygote
         exclude:
           - version: '1.6'
             group: Formalities

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/DifferentiationInterfaceForwardDiffExt.jl

@@ -12,7 +12,8 @@ using DifferentiationInterface:
     NoDerivativeExtras,
     NoSecondDerivativeExtras,
     PushforwardExtras
-using ForwardDiff.DiffResults: DiffResults, DiffResult, GradientResult, MutableDiffResult
+using ForwardDiff.DiffResults:
+    DiffResults, DiffResult, GradientResult, HessianResult, MutableDiffResult
 using ForwardDiff:
     Chunk,
     Dual,

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/onearg.jl

@@ -132,10 +132,11 @@ function DI.value_and_gradient!(
 end
 
 function DI.value_and_gradient(
-    f::F, backend::AutoForwardDiff, x, extras::ForwardDiffGradientExtras
+    f::F, ::AutoForwardDiff, x, extras::ForwardDiffGradientExtras
 ) where {F}
-    grad = similar(x)
-    return DI.value_and_gradient!(f, grad, backend, x, extras)
+    result = GradientResult(x)
+    result = gradient!(result, f, x, extras.config)
+    return DiffResults.value(result), DiffResults.gradient(result)
 end
 
 function DI.gradient!(
@@ -189,19 +190,22 @@ end
 
 ## Hessian
 
-struct ForwardDiffHessianExtras{C1,C2} <: HessianExtras
+struct ForwardDiffHessianExtras{C1,C2,C3} <: HessianExtras
     array_config::C1
-    result_config::C2
+    manual_result_config::C2
+    auto_result_config::C3
 end
 
 function DI.prepare_hessian(f, backend::AutoForwardDiff, x)
-    example_result = MutableDiffResult(
+    manual_result = MutableDiffResult(
         one(eltype(x)), (similar(x), similar(x, length(x), length(x)))
     )
+    auto_result = HessianResult(x)
     chunk = choose_chunk(backend, x)
     array_config = HessianConfig(f, x, chunk)
-    result_config = HessianConfig(f, example_result, x, chunk)
-    return ForwardDiffHessianExtras(array_config, result_config)
+    manual_result_config = HessianConfig(f, manual_result, x, chunk)
+    auto_result_config = HessianConfig(f, auto_result, x, chunk)
+    return ForwardDiffHessianExtras(array_config, manual_result_config, auto_result_config)
 end
 
 function DI.hessian!(
@@ -218,7 +222,7 @@ function DI.value_gradient_and_hessian!(
     f::F, grad, hess, ::AutoForwardDiff, x, extras::ForwardDiffHessianExtras
 ) where {F}
     result = MutableDiffResult(one(eltype(x)), (grad, hess))
-    result = hessian!(result, f, x, extras.result_config)
+    result = hessian!(result, f, x, extras.manual_result_config)
     return (
         DiffResults.value(result), DiffResults.gradient(result), DiffResults.hessian(result)
     )
@@ -227,10 +231,8 @@ end
 function DI.value_gradient_and_hessian(
     f::F, ::AutoForwardDiff, x, extras::ForwardDiffHessianExtras
 ) where {F}
-    result = MutableDiffResult(
-        one(eltype(x)), (similar(x), similar(x, length(x), length(x)))
-    )
-    result = hessian!(result, f, x, extras.result_config)
+    result = HessianResult(x)
+    result = hessian!(result, f, x, extras.auto_result_config)
     return (
         DiffResults.value(result), DiffResults.gradient(result), DiffResults.hessian(result)
     )

DifferentiationInterface/test/Double/ChainRulesCore-Zygote/test.jl

@@ -11,7 +11,7 @@ end
 
 test_differentiation(
     AutoChainRules(ZygoteRuleConfig());
-    excluded=[SecondDerivativeScenario],
+    excluded=[:second_derivative],
     second_order=VERSION >= v"1.10",
     logging=LOGGING,
 );

DifferentiationInterface/test/Single/Enzyme/test.jl

@@ -48,14 +48,14 @@ test_differentiation(
         SecondOrder(AutoEnzyme(; mode=Enzyme.Forward), AutoEnzyme(; mode=Enzyme.Reverse)),
     ];
     first_order=false,
-    excluded=[SecondDerivativeScenario],
+    excluded=[:second_derivative],
     logging=LOGGING,
 );
 
 test_differentiation(
     [AutoEnzyme(; mode=nothing), AutoEnzyme(; mode=Enzyme.Forward)];
     first_order=false,
-    excluded=[HessianScenario, HVPScenario],
+    excluded=[:hessian, :hvp],
     logging=LOGGING,
 );
 
@@ -72,7 +72,7 @@ test_differentiation(
 test_differentiation(
     sparse_backends,
     default_scenarios();
-    excluded=[DerivativeScenario, GradientScenario, PullbackScenario, PushforwardScenario],
+    excluded=[:derivative, :gradient, :pullback, :pushforward],
     second_order=false,
     logging=LOGGING,
 );

DifferentiationInterface/test/Single/FiniteDiff/test.jl

@@ -8,6 +8,4 @@ for backend in [AutoFiniteDiff()]
     @test check_hessian(backend)
 end
 
-test_differentiation(
-    AutoFiniteDiff(); excluded=[SecondDerivativeScenario, HVPScenario], logging=LOGGING
-);
+test_differentiation(AutoFiniteDiff(); excluded=[:second_derivative, :hvp], logging=LOGGING);

DifferentiationInterface/test/Single/ForwardDiff/efficiency.jl

@@ -36,7 +36,7 @@ test_differentiation(
     # ForwardDiff accesses individual indices
     vcat(component_scenarios(), static_scenarios());
     # jacobian is super slow for some reason
-    excluded=[JacobianScenario],
+    excluded=[:jacobian],
     second_order=false,
     logging=LOGGING,
 );
@@ -46,14 +46,7 @@ test_differentiation(
 test_differentiation(
     sparse_backends,
     default_scenarios();
-    excluded=[
-        DerivativeScenario,
-        GradientScenario,
-        HVPScenario,
-        PullbackScenario,
-        PushforwardScenario,
-        SecondDerivativeScenario,
-    ],
+    excluded=[:derivative, :gradient, :hvp, :pullback, :pushforward, :second_derivative],
     logging=LOGGING,
 );
 

DifferentiationInterface/test/Single/ForwardDiff/test.jl

@@ -21,32 +21,23 @@ end
 
 ## Dense backends
 
-test_differentiation(AutoZygote(); excluded=[SecondDerivativeScenario], logging=LOGGING);
-
-test_differentiation(
-    AutoZygote(),
-    vcat(component_scenarios(), static_scenarios());
-    second_order=false,
-    logging=LOGGING,
-)
+test_differentiation(AutoZygote(); excluded=[:second_derivative], logging=LOGGING);
 
 if VERSION >= v"1.10"
-    test_differentiation(AutoZygote(), gpu_scenarios(); second_order=false, logging=LOGGING)
+    test_differentiation(
+        AutoZygote(),
+        vcat(component_scenarios(), gpu_scenarios(), static_scenarios());
+        second_order=false,
+        logging=LOGGING,
+    )
 end
 
 ## Sparse backends
 
 test_differentiation(
     sparse_backends,
     default_scenarios();
-    excluded=[
-        DerivativeScenario,
-        GradientScenario,
-        HVPScenario,
-        PullbackScenario,
-        PushforwardScenario,
-        SecondDerivativeScenario,
-    ],
+    excluded=[:derivative, :gradient, :hvp, :pullback, :pushforward, :second_derivative],
     logging=LOGGING,
 );
 

DifferentiationInterface/test/Single/Zygote/test.jl

@@ -1,13 +1,14 @@
 name = "DifferentiationInterfaceTest"
 uuid = "a82114a7-5aa3-49a8-9643-716bb13727a3"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.4.4"
+version = "0.5.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
+DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
@@ -24,6 +25,7 @@ ADTypes = "1.0.0"
 Chairmarks = "1.2.1"
 Compat = "4"
 ComponentArrays = "0.15"
+DataFrames = "1.6.1"
 DifferentiationInterface = "0.5.4"
 DocStringExtensions = "0.9"
 JET = "0.4 - 0.8, 0.9"
@@ -50,6 +52,7 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
 SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["ADTypes", "Aqua", "DataFrames", "DifferentiationInterface", "ForwardDiff", "JET", "JuliaFormatter", "Pkg", "SparseArrays", "SparseConnectivityTracer", "SparseMatrixColorings", "Test"]
+test = ["ADTypes", "Aqua", "DataFrames", "DifferentiationInterface", "ForwardDiff", "JET", "JuliaFormatter", "Pkg", "SparseArrays", "SparseConnectivityTracer", "SparseMatrixColorings", "Test", "Zygote"]

DifferentiationInterfaceTest/Project.toml

@@ -1,10 +1,10 @@
+# API reference
+
 ```@meta
 CurrentModule = Main
 CollapsedDocStrings = true
 ```
 
-# API reference
-
 ```@docs
 DifferentiationInterfaceTest
 ```
@@ -30,7 +30,6 @@ static_scenarios
 ## Scenario types
 
 ```@docs
-AbstractScenario
 PushforwardScenario
 PullbackScenario
 DerivativeScenario

DifferentiationInterfaceTest/docs/src/api.md

@@ -5,7 +5,7 @@ We present a typical workflow with DifferentiationInterfaceTest.jl, building on
 ```@repl tuto
 using DifferentiationInterface, DifferentiationInterfaceTest
 import ForwardDiff, Enzyme
-import DataFrames, Markdown, PrettyTables, Printf
+import Markdown, PrettyTables, Printf
 ```
 
 ## Introduction
@@ -31,15 +31,18 @@ Of course we know the true gradient mapping:
 DifferentiationInterfaceTest.jl relies with so-called "scenarios", in which you encapsulate the information needed for your test:
 
 - the function `f`
-- the input `x` (and output `y` for mutating functions)
-- optionally a reference `ref` to check against
+- the input `x` and output `y`
+- the number of arguments for `f` (either `1` or `2`)
+- the behavior of the operator (either `:inplace` or `:outofplace`)
 
-There is one scenario per operator, and so here we will use [`GradientScenario`](@ref):
+There is one scenario constructor per operator, and so here we will use [`GradientScenario`](@ref):
 
 ```@example tuto
+xv = rand(Float32, 3)
+xm = rand(Float64, 3, 2)
 scenarios = [
-    GradientScenario(f; x=rand(Float32, 3), ref=∇f, place=:inplace),
-    GradientScenario(f; x=rand(Float64, 3, 2), ref=∇f, place=:inplace)
+    GradientScenario(f; x=xv, y=f(xv), nb_args=1, place=:inplace),
+    GradientScenario(f; x=xm, y=f(xm), nb_args=1, place=:inplace)
 ];
 nothing  # hide
 ```
@@ -67,17 +70,11 @@ Once you are confident that your backends give the correct answers, you probably
 This is made easy by the [`benchmark_differentiation`](@ref) function, whose syntax should feel familiar:
 
 ```@example tuto
-benchmark_result = benchmark_differentiation(backends, scenarios);
+df = benchmark_differentiation(backends, scenarios);
 ```
 
-The resulting object is a `Vector` of [`DifferentiationBenchmarkDataRow`](@ref), which can easily be converted into a `DataFrame` from [DataFrames.jl](https://github.com/JuliaData/DataFrames.jl):
-
-```@example tuto
-df = DataFrames.DataFrame(benchmark_result)
-```
-
-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.
+The resulting object is `DataFrame` from [DataFrames.jl](https://github.com/JuliaData/DataFrames.jl), whose columns correspond to the fields of [`DifferentiationBenchmarkDataRow`](@ref):
+Here's what it looks like with all of its columns.
 
 ```@example tuto
 table = PrettyTables.pretty_table(