[BREAKING] Clean up tests (#470)

* Clean up tests

* Fixes

* Fix sparse autodiff

* Remove options

* Typo

* Remove exports

* Fixes

* Docs

* Fixes

* Fix

* Fix order of scenario modificators

Author: gdalle

.github/workflows/Test.yml

@@ -4,7 +4,7 @@ on:
   push:
     branches:
       - main
-    tags: ['*']
+    tags: ["*"]
   pull_request:
     types: [opened, reopened, synchronize, ready_for_review]
   workflow_dispatch:
@@ -16,7 +16,6 @@ concurrency:
   cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}
 
 jobs:
-  
   test-DI:
     name: ${{ matrix.version }} - DI (${{ matrix.group }})
     runs-on: ubuntu-latest
@@ -28,9 +27,9 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1'
-          - 'lts'
-          - 'pre'
+          - "1"
+          - "lts"
+          - "pre"
         group:
           - Formalities
           - Internals
@@ -48,51 +47,53 @@ jobs:
           - Back/Tapir
           - Back/Tracker
           - Back/Zygote
-          - Down/Detector
-          - Down/DifferentiateWith
+          - Misc/DifferentiateWith
+          - Misc/FromPrimitive
+          - Misc/SparsityDetector
+          - Misc/ZeroBackends
           - Down/Flux
           - Down/Lux
         exclude:
           # lts
-          - version: 'lts'
+          - version: "lts"
             group: Formalities
-          - version: 'lts'
+          - version: "lts"
             group: Back/ChainRulesCore
-          - version: 'lts'
+          - version: "lts"
             group: Back/Diffractor
-          - version: 'lts'
+          - version: "lts"
             group: Back/Enzyme
-          - version: 'lts'
+          - version: "lts"
             group: Back/FiniteDiff
-          - version: 'lts'
+          - version: "lts"
             group: Back/FastDifferentiation
-          - version: 'lts'
+          - version: "lts"
             group: Back/PolyesterForwardDiff
-          - version: 'lts'
+          - version: "lts"
             group: Back/SecondOrder
-          - version: 'lts'
+          - version: "lts"
             group: Back/Symbolics
-          - version: 'lts'
+          - version: "lts"
             group: Back/Tapir
-          - version: 'lts'
-            group: Down/Detector
-          - version: 'lts'
+          - version: "lts"
+            group: Misc/SparsityDetector
+          - version: "lts"
             group: Down/Flux
-          - version: 'lts'
+          - version: "lts"
             group: Down/Lux
           # pre-release
-          - version: 'pre'
+          - version: "pre"
             group: Formalities
-          - version: 'pre'
+          - version: "pre"
             group: Back/ChainRulesCore
-          - version: 'pre'
+          - version: "pre"
             group: Back/Enzyme
-          - version: 'pre'
+          - version: "pre"
             group: Back/Tapir
-          - version: 'pre'
+          - version: "pre"
             group: Back/SecondOrder
-          - version: 'pre'
-            group: Down/Detector
+          - version: "pre"
+            group: Misc/SparsityDetector
     env:
       SHOULDRUN: ${{ matrix.version == '1' || !github.event.pull_request.draft }}
       JULIA_DI_TEST_GROUP: ${{ matrix.group }}
@@ -138,18 +139,18 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1'
-          - 'lts'
-          - 'pre'
+          - "1"
+          - "lts"
+          - "pre"
         group:
           - Formalities
           - Zero
-          - ForwardDiff
+          - Standard
           - Weird
         exclude:
-          - version: 'lts'
+          - version: "lts"
             group: Formalities
-          - version: 'lts'
+          - version: "lts"
             group: Weird
     env:
       SHOULDRUN: ${{ matrix.version == '1' || !github.event.pull_request.draft }}
@@ -183,4 +184,4 @@ jobs:
           flags: DIT
           name: ${{ matrix.version }} - DIT (${{ matrix.group }})
           token: ${{ secrets.CODECOV_TOKEN }}
-          fail_ci_if_error: true
+          fail_ci_if_error: true

DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/DifferentiationInterfaceSparseMatrixColoringsExt.jl

@@ -41,6 +41,12 @@ using SparseMatrixColorings:
     decompress,
     decompress!
 
+with_context(f, contexts::Vararg{Context,C}) where {C} = (DI.with_context(f, contexts...),)
+
+function with_context(f!, y, contexts::Vararg{Context,C}) where {C}
+    return (DI.with_context(f!, contexts...), y)
+end
+
 include("jacobian.jl")
 include("hessian.jl")
 

DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/hessian.jl

@@ -39,7 +39,9 @@ function DI.prepare_hessian(
     f::F, backend::AutoSparse, x, contexts::Vararg{Context,C}
 ) where {F,C}
     dense_backend = dense_ad(backend)
-    sparsity = hessian_sparsity(f, x, sparsity_detector(backend))
+    sparsity = hessian_sparsity(
+        with_context(f, contexts...)..., x, sparsity_detector(backend)
+    )
     problem = ColoringProblem{:symmetric,:column}()
     coloring_result = coloring(
         sparsity, problem, coloring_algorithm(backend); decompression_eltype=eltype(x)

DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian.jl

@@ -81,7 +81,10 @@ function _prepare_sparse_jacobian_aux(
     ::PushforwardFast, y, f_or_f!y::FY, backend::AutoSparse, x, contexts::Vararg{Context,C}
 ) where {FY,C}
     dense_backend = dense_ad(backend)
-    sparsity = jacobian_sparsity(f_or_f!y..., x, sparsity_detector(backend))
+
+    sparsity = jacobian_sparsity(
+        with_context(f_or_f!y..., contexts...)..., x, sparsity_detector(backend)
+    )
     problem = ColoringProblem{:nonsymmetric,:column}()
     coloring_result = coloring(
         sparsity,
@@ -115,7 +118,9 @@ function _prepare_sparse_jacobian_aux(
     ::PushforwardSlow, y, f_or_f!y::FY, backend::AutoSparse, x, contexts::Vararg{Context,C}
 ) where {FY,C}
     dense_backend = dense_ad(backend)
-    sparsity = jacobian_sparsity(f_or_f!y..., x, sparsity_detector(backend))
+    sparsity = jacobian_sparsity(
+        with_context(f_or_f!y..., contexts...)..., x, sparsity_detector(backend)
+    )
     problem = ColoringProblem{:nonsymmetric,:row}()
     coloring_result = coloring(
         sparsity,
@@ -163,7 +168,7 @@ end
 function DI.value_and_jacobian(
     f::F, extras::SparseJacobianExtras, backend::AutoSparse, x, contexts::Vararg{Context,C}
 ) where {F,C}
-    return f(x), jacobian(f, extras, backend, x, contexts...)
+    return f(x, map(unwrap, contexts)...), jacobian(f, extras, backend, x, contexts...)
 end
 
 function DI.value_and_jacobian!(

DifferentiationInterface/test/Back/Enzyme/test.jl

@@ -55,7 +55,7 @@ test_differentiation(
 
 test_differentiation(
     duplicated_function_backends,
-    DIT.make_closure.(default_scenarios());
+    default_scenarios(; include_normal=false, include_closurified=true);
     second_order=false,
     logging=LOGGING,
 );
@@ -81,7 +81,7 @@ test_differentiation(
 
 test_differentiation(
     AutoEnzyme(; mode=Enzyme.Forward),  # TODO: add more
-    DIT.remove_batched(default_scenarios());
+    default_scenarios(; include_batchified=false);
     correctness=false,
     type_stability=true,
     second_order=false,

DifferentiationInterface/test/Back/ForwardDiff/test.jl

@@ -31,7 +31,7 @@ test_differentiation(dense_backends, default_scenarios(); logging=LOGGING);
 
 test_differentiation(
     dense_backends,
-    default_scenarios();
+    default_scenarios(; include_constantified=true);
     correctness=false,
     type_stability=true,
     second_order=false,
@@ -57,4 +57,9 @@ test_differentiation(
     logging=LOGGING,
 );
 
-test_differentiation(sparse_backends, sparse_scenarios(); sparsity=true, logging=LOGGING);
+test_differentiation(
+    sparse_backends,
+    sparse_scenarios(; include_constantified=true);
+    sparsity=true,
+    logging=LOGGING,
+);

…face/test/Down/DifferentiateWith/test.jl …face/test/Misc/DifferentiateWith/test.jlDifferentiationInterface/test/Down/DifferentiateWith/test.jl renamed to DifferentiationInterface/test/Misc/DifferentiateWith/test.jl DifferentiationInterface/test/Down/DifferentiateWith/test.jl renamed to DifferentiationInterface/test/Misc/DifferentiateWith/test.jl

@@ -10,10 +10,10 @@ using Test
 LOGGING = get(ENV, "CI", "false") == "false"
 
 function zygote_breaking_scenarios()
-    onearg_scens = filter(default_scenarios()) do scen
-        DIT.nb_args(scen) == 1
+    outofplace_scens = filter(default_scenarios()) do scen
+        DIT.operator_place(scen) == :out
     end
-    bad_onearg_scens = map(onearg_scens) do scen
+    bad_outofplace_scens = map(outofplace_scens) do scen
         function bad_f(x)
             a = Vector{eltype(x)}(undef, 1)
             a[1] = sum(x)
@@ -23,7 +23,7 @@ function zygote_breaking_scenarios()
         bad_scen = DIT.change_function(scen, wrapped_bad_f)
         return bad_scen
     end
-    return bad_onearg_scens
+    return bad_outofplace_scens
 end
 
 test_differentiation(

…terface/test/Internals/from_primitive.jl …nterface/test/Misc/FromPrimitive/test.jlDifferentiationInterface/test/Internals/from_primitive.jl renamed to DifferentiationInterface/test/Misc/FromPrimitive/test.jl DifferentiationInterface/test/Internals/from_primitive.jl renamed to DifferentiationInterface/test/Misc/FromPrimitive/test.jl

@@ -1,7 +1,6 @@
 using DifferentiationInterface, DifferentiationInterfaceTest
 using DifferentiationInterface: AutoForwardFromPrimitive, AutoReverseFromPrimitive
 using DifferentiationInterfaceTest
-using DifferentiationInterfaceTest: insert_context
 using ForwardDiff: ForwardDiff
 using Test
 
@@ -17,8 +16,6 @@ for backend in vcat(fromprimitive_backends)
     @test check_inplace(backend)
 end
 
-test_differentiation(fromprimitive_backends, default_scenarios(); logging=LOGGING);
-
 test_differentiation(
-    fromprimitive_backends, insert_context.(default_scenarios()); logging=LOGGING
+    fromprimitive_backends, default_scenarios(; include_constantified=true); logging=LOGGING
 );

…tionInterface/test/Down/Detector/test.jl …rface/test/Misc/SparsityDetector/test.jlDifferentiationInterface/test/Down/Detector/test.jl renamed to DifferentiationInterface/test/Misc/SparsityDetector/test.jl DifferentiationInterface/test/Down/Detector/test.jl renamed to DifferentiationInterface/test/Misc/SparsityDetector/test.jl

@@ -1,7 +1,6 @@
 using DifferentiationInterface
 using DifferentiationInterface: AutoZeroForward, AutoZeroReverse
 using DifferentiationInterfaceTest
-using DifferentiationInterfaceTest: insert_context
 using ComponentArrays: ComponentArrays
 using JLArrays: JLArrays
 using StaticArrays: StaticArrays
@@ -20,7 +19,7 @@ end
 
 test_differentiation(
     zero_backends,
-    zero.(default_scenarios());
+    zero.(default_scenarios(; include_constantified=true));
     correctness=true,
     type_stability=true,
     excluded=[:second_derivative],
@@ -39,15 +38,6 @@ test_differentiation(
     logging=LOGGING,
 )
 
-## Contexts
-
-test_differentiation(
-    zero_backends,
-    insert_context.(zero.(default_scenarios()));
-    correctness=true,
-    logging=LOGGING,
-)
-
 ## Weird arrays
 
 test_differentiation(