[BREAKING] Improve type stability tests and benchmarking (#560)

* Improve type stability tests and benchmarking

* Remove `first_order` and `second_order`

* Docs

* Zero allocs

* Fixes

* Call count

* Fix

* Fix

* Add count calls

* Default count calls

* Fix

Author: gdalle

.github/workflows/Test.yml

@@ -25,7 +25,7 @@ jobs:
       actions: write
       contents: read
     strategy:
-      fail-fast: true
+      fail-fast: false  # TODO: toggle
       matrix:
         version:
           - "1.10"

DifferentiationInterface/test/Back/ChainRulesBackends/chainrules_zygote.jl

@@ -23,6 +23,6 @@ test_differentiation(
 test_differentiation(
     AutoChainRules(ZygoteRuleConfig()),
     default_scenarios(; include_normal=false, include_constantified=true);
-    second_order=false,
+    excluded=SECOND_ORDER,
     logging=LOGGING,
 );

DifferentiationInterface/test/Back/ChainRulesBackends/diffractor.jl

@@ -13,5 +13,8 @@ for backend in [AutoDiffractor()]
 end
 
 test_differentiation(
-    AutoDiffractor(), default_scenarios(; linalg=false); second_order=false, logging=LOGGING
+    AutoDiffractor(),
+    default_scenarios(; linalg=false);
+    excluded=SECOND_ORDER,
+    logging=LOGGING,
 );

DifferentiationInterface/test/Back/Enzyme/test.jl

@@ -31,19 +31,19 @@ end;
 
 ## First order
 
-test_differentiation(backends, default_scenarios(); second_order=false, logging=LOGGING);
+test_differentiation(backends, default_scenarios(); excluded=SECOND_ORDER, logging=LOGGING);
 
 test_differentiation(
     backends[1:3],
     default_scenarios(; include_normal=false, include_constantified=true);
-    second_order=false,
+    excluded=SECOND_ORDER,
     logging=LOGGING,
 );
 
 test_differentiation(
     duplicated_backends,
     default_scenarios(; include_normal=false, include_closurified=true);
-    second_order=false,
+    excluded=SECOND_ORDER,
     logging=LOGGING,
 );
 
@@ -54,8 +54,8 @@ test_differentiation(
     AutoEnzyme(; mode=Enzyme.Forward),  # TODO: add more
     default_scenarios(; include_batchified=false);
     correctness=false,
-    type_stability=true,
-    second_order=false,
+    type_stability=:prepared,
+    excluded=SECOND_ORDER,
     logging=LOGGING,
 );
 =#
@@ -65,27 +65,24 @@ test_differentiation(
 test_differentiation(
     AutoEnzyme(),
     default_scenarios(; include_constantified=true);
-    first_order=false,
+    excluded=FIRST_ORDER,
     logging=LOGGING,
 );
 
 test_differentiation(
     AutoEnzyme(; mode=Enzyme.Forward);
-    first_order=false,
-    excluded=[:hessian, :hvp],
+    excluded=vcat(FIRST_ORDER, [:hessian, :hvp]),
     logging=LOGGING,
 );
 
 test_differentiation(
     AutoEnzyme(; mode=Enzyme.Reverse);
-    first_order=false,
-    excluded=[:second_derivative],
+    excluded=vcat(FIRST_ORDER, [:second_derivative]),
     logging=LOGGING,
 );
 
 test_differentiation(
     SecondOrder(AutoEnzyme(; mode=Enzyme.Reverse), AutoEnzyme(; mode=Enzyme.Forward));
-    first_order=false,
     logging=LOGGING,
 );
 

DifferentiationInterface/test/Back/FiniteDifferences/test.jl

@@ -15,6 +15,6 @@ end
 test_differentiation(
     AutoFiniteDifferences(; fdm=FiniteDifferences.central_fdm(3, 1)),
     default_scenarios(; include_constantified=true);
-    second_order=false,
+    excluded=SECOND_ORDER,
     logging=LOGGING,
 );

DifferentiationInterface/test/Back/ForwardDiff/test.jl

@@ -23,21 +23,21 @@ test_differentiation(
 );
 
 test_differentiation(
-    AutoForwardDiff(); correctness=false, type_stability=true, logging=LOGGING
+    AutoForwardDiff(); correctness=false, type_stability=:prepared, logging=LOGGING
 );
 
 test_differentiation(
     AutoForwardDiff(; chunksize=5);
     correctness=false,
-    type_stability=(; preparation=true, prepared_op=true, unprepared_op=false),
+    type_stability=:full,
+    excluded=[:hessian],
     logging=LOGGING,
 );
 
 test_differentiation(
     backends,
     vcat(component_scenarios(), static_scenarios()); # FD accesses individual indices
-    excluded=[:jacobian],  # jacobian is super slow for some reason
-    second_order=false,
+    excluded=vcat(SECOND_ORDER, [:jacobian]),  # jacobian is super slow for some reason
     logging=LOGGING,
 );
 

DifferentiationInterface/test/Back/Mooncake/test.jl

@@ -17,6 +17,6 @@ end
 test_differentiation(
     backends,
     default_scenarios(; include_constantified=true);
-    second_order=false,
+    excluded=SECOND_ORDER,
     logging=LOGGING,
 );

DifferentiationInterface/test/Back/Tracker/test.jl

@@ -15,6 +15,6 @@ end
 test_differentiation(
     AutoTracker(),
     default_scenarios(; include_constantified=true);
-    second_order=false,
+    excluded=SECOND_ORDER,
     logging=LOGGING,
 );

DifferentiationInterface/test/Back/Zygote/test.jl

@@ -28,12 +28,12 @@ test_differentiation(
     logging=LOGGING,
 );
 
-test_differentiation(second_order_backends; first_order=false, logging=LOGGING);
+test_differentiation(second_order_backends; logging=LOGGING);
 
 test_differentiation(
     backends[1],
     vcat(component_scenarios(), gpu_scenarios());
-    second_order=false,
+    excluded=SECOND_ORDER,
     logging=LOGGING,
 )
 

DifferentiationInterface/test/Misc/DifferentiateWith/test.jl

@@ -24,6 +24,6 @@ end
 test_differentiation(
     [AutoForwardDiff(), AutoZygote()],
     differentiatewith_scenarios();
-    second_order=false,
+    excluded=SECOND_ORDER,
     logging=LOGGING,
 )