Call count in benchmark (#151)

* Call count in benchmark

* Fix

Author: gdalle

DifferentiationInterface/ext/DifferentiationInterfaceFastDifferentiationExt/DifferentiationInterfaceFastDifferentiationExt.jl

@@ -31,6 +31,8 @@ const AnyAutoFastDifferentiation = Union{
 }
 
 DI.mode(::AnyAutoFastDifferentiation) = ADTypes.AbstractSymbolicDifferentiationMode
+DI.pushforward_performance(::AnyAutoFastDifferentiation) = DI.PushforwardFast()
+DI.pullback_performance(::AnyAutoFastDifferentiation) = DI.PullbackSlow()
 
 myvec(x::Number) = [x]
 myvec(x::AbstractArray) = vec(x)

DifferentiationInterface/src/pullback.jl

@@ -29,9 +29,12 @@ prepare_pullback(f!, ::AbstractADType, y, x) = NoPullbackExtras()
 function value_and_pullback_split(
     f, backend::AbstractADType, x, extras::PullbackExtras=prepare_pullback(f, backend, x)
 )
-    if !Bool(pushforward_performance(backend))
-        error("Pushforward not available for backend $backend")
-    end
+    return value_and_pullback_split_aux(
+        f, backend, x, extras, pullback_performance(backend)
+    )
+end
+
+function value_and_pullback_split_aux(f, backend, x, extras, ::PullbackSlow)
     pushforward_extras = prepare_pushforward(f, backend, x)
     y = f(x)
     pullbackfunc = if x isa Number && y isa Number

DifferentiationInterface/src/pushforward.jl

@@ -33,9 +33,12 @@ function value_and_pushforward(
     dx,
     extras::PushforwardExtras=prepare_pushforward(f, backend, x),
 )
-    if !Bool(pullback_performance(backend))
-        error("Pullback not available for backend $backend")
-    end
+    return value_and_pushforward_aux(
+        f, backend, x, dx, extras, pushforward_performance(backend)
+    )
+end
+
+function value_and_pushforward_aux(f, backend, x, dx, extras, ::PushforwardSlow)
     pullback_extras = prepare_pullback(f, backend, x)
     y, pullbackfunc = value_and_pullback_split(f, backend, x, pullback_extras)
     dy = if x isa Number && y isa Number

DifferentiationInterfaceTest/docs/src/tutorial.md

@@ -45,7 +45,6 @@ Let us experiment with various input types and sizes:
 scenarios = [
     GradientScenario(f; x=rand(Float64, 3), ref=∇f),
     GradientScenario(f; x=rand(Float32, 3, 4), ref=∇f),
-    GradientScenario(f; x=rand(Float16, 3, 4, 5), ref=∇f),
 ];
 ```
 

DifferentiationInterfaceTest/src/DifferentiationInterfaceTest.jl

@@ -52,7 +52,6 @@ include("utils/filter.jl")
 
 include("tests/correctness.jl")
 include("tests/type_stability.jl")
-include("tests/call_count.jl")
 include("tests/sparsity.jl")
 include("tests/benchmark.jl")
 include("test_differentiation.jl")

DifferentiationInterfaceTest/src/test_differentiation.jl

@@ -14,7 +14,6 @@ Testing:
 - `correctness=true`: whether to compare the differentiation results with the theoretical values specified in each scenario
     - If a backend object like `correctness=AutoForwardDiff()` is passed instead of a boolean, the results will be compared using that reference backend as the ground truth.
     - Otherwise, the scenario-specific reference operator will be used as the ground truth instead, see [`AbstractScenario`](@ref) for details.
-- `call_count=false`: whether to check that the function is called the right number of times
 - `type_stability=false`: whether to check type stability with JET.jl (thanks to `@test_opt`)
 - `sparsity`: whether to check sparsity of the jacobian / hessian
 - `detailed=false`: whether to print a detailed or condensed test log
@@ -79,13 +78,18 @@ function test_differentiation(
     prog = ProgressUnknown(; desc="$title", spinner=true, enabled=logging)
 
     @testset verbose = true "$title" begin
-        @testset verbose = detailed "$(backend_string(backend))" for backend in backends
-            @testset "$scen" for scen in filter(s -> compatible(backend, s), scenarios)
+        @testset verbose = detailed "$(backend_string(backend))" for (i, backend) in
+                                                                     enumerate(backends)
+            filtered_scenarios = filter(s -> compatible(backend, s), scenarios)
+            @testset "$scen" for (j, scen) in enumerate(filtered_scenarios)
                 next!(
                     prog;
                     showvalues=[
-                        (:backend, backend_string(backend)),
-                        (:scenario, typeof(scen).name.name),
+                        (:backend, "$(backend_string(backend)) - $i/$(length(backends))"),
+                        (
+                            :scenario,
+                            "$(typeof(scen).name.name) - $j/$(length(filtered_scenarios))",
+                        ),
                         (:function, scen.f),
                         (:input, typeof(scen.x)),
                         (:output, typeof(scen.y)),
@@ -102,9 +106,6 @@ function test_differentiation(
                         )
                     end
                 end
-                call_count && @testset "Call count" begin
-                    test_call_count(backend, scen)
-                end
                 type_stability && @testset "Type stability" begin
                     test_jet(backend, scen)
                 end
@@ -163,13 +164,17 @@ function benchmark_differentiation(
 
     benchmark_data = BenchmarkDataRow[]
     prog = ProgressUnknown(; desc="Benchmarking", spinner=true, enabled=logging)
-    for backend in backends
-        for scen in filter(s -> compatible(backend, s), scenarios)
+    for (i, backend) in enumerate(backends)
+        filtered_scenarios = filter(s -> compatible(backend, s), scenarios)
+        for (j, scen) in enumerate(filtered_scenarios)
             next!(
                 prog;
                 showvalues=[
-                    (:backend, backend_string(backend)),
-                    (:scenario, typeof(scen).name.name),
+                    (:backend, "$(backend_string(backend)) - $i/$(length(backends))"),
+                    (
+                        :scenario,
+                        "$(typeof(scen).name.name) - $j/$(length(filtered_scenarios))",
+                    ),
                     (:function, scen.f),
                     (:input, typeof(scen.x)),
                     (:output, typeof(scen.y)),