Make CI faster

.github/workflows/Test.yml

@@ -33,7 +33,7 @@ jobs:
         group:
           - Misc/Internals
           - Misc/DifferentiateWith
-          - Misc/FromPrimitive
+          - Misc/SimpleFiniteDiff
           - Misc/SparsityDetector
           - Misc/ZeroBackends
           - Back/ChainRules

DifferentiationInterface/src/DifferentiationInterface.jl

@@ -62,6 +62,7 @@ include("fallbacks/change_prep.jl")
 include("misc/differentiate_with.jl")
 include("misc/from_primitive.jl")
 include("misc/sparsity_detector.jl")
+include("misc/simple_finite_diff.jl")
 include("misc/zero_backends.jl")
 
 ## Exported

DifferentiationInterface/src/misc/from_primitive.jl

@@ -19,8 +19,9 @@ function BatchSizeSettings(fromprim::FromPrimitive, N::Integer)
     return BatchSizeSettings(fromprim.backend, N)
 end
 
-## Forward
+## Forward (no longer used)
 
+#=
 struct AutoForwardFromPrimitive{B} <: FromPrimitive
     backend::B
 end
@@ -104,6 +105,7 @@ function value_and_pushforward!(
         f!, y, ty, prep.pushforward_prep, fromprim.backend, x, tx, contexts...
     )
 end
+=#
 
 ## Reverse
 

DifferentiationInterface/src/misc/simple_finite_diff.jl

@@ -0,0 +1,93 @@
+"""
+    AutoSimpleFiniteDiff <: ADTypes.AbstractADType
+
+Forward mode backend based on the finite difference `(f(x + ε) - f(x)) / ε`, with artificial chunk size to mimick ForwardDiff.
+
+# Constructor
+
+    AutoSimpleFiniteDiff(ε=1e-5; chunksize=nothing)
+"""
+struct AutoSimpleFiniteDiff{chunksize,T<:Real} <: AbstractADType
+    ε::T
+end
+
+function AutoSimpleFiniteDiff(ε=1e-5; chunksize=nothing)
+    return AutoSimpleFiniteDiff{chunksize,typeof(ε)}(ε)
+end
+
+ADTypes.mode(::AutoSimpleFiniteDiff) = ForwardMode()
+check_available(::AutoSimpleFiniteDiff) = true
+inplace_support(::AutoSimpleFiniteDiff) = InPlaceSupported()
+
+function BatchSizeSettings(::AutoSimpleFiniteDiff{nothing}, N::Integer)
+    B = min(12, N)
+    singlebatch = B == N
+    aligned = N % B == 0
+    return BatchSizeSettings{B,singlebatch,aligned}(N)
+end
+
+function BatchSizeSettings(::AutoSimpleFiniteDiff{chunksize}, N::Integer) where {chunksize}
+    @assert chunksize <= N
+    B = chunksize
+    singlebatch = B == N
+    aligned = N % B == 0
+    return BatchSizeSettings{B,singlebatch,aligned}(N)
+end
+
+function threshold_batchsize(
+    backend::AutoSimpleFiniteDiff{chunksize1}, chunksize2::Integer
+) where {chunksize1}
+    chunksize = isnothing(chunksize1) ? nothing : min(chunksize1, chunksize2)
+    return AutoSimpleFiniteDiff(backend.ε; chunksize)
+end
+
+function prepare_pushforward(
+    f::F, ::AutoSimpleFiniteDiff, x, tx::NTuple, contexts::Vararg{Context,C}
+) where {F,C}
+    return NoPushforwardPrep()
+end
+
+function prepare_pushforward(
+    f!::F, y, ::AutoSimpleFiniteDiff, x, tx::NTuple, contexts::Vararg{Context,C}
+) where {F,C}
+    return NoPushforwardPrep()
+end
+
+function value_and_pushforward(
+    f::F,
+    ::NoPushforwardPrep,
+    backend::AutoSimpleFiniteDiff,
+    x,
+    tx::NTuple{B},
+    contexts::Vararg{Context,C},
+) where {F,B,C}
+    ε = eltype(x)(backend.ε)
+    y = f(x, map(unwrap, contexts)...)
+    ty = map(tx) do dx
+        yε = f(x + ε * dx, map(unwrap, contexts)...)
+        y0 = f(x, map(unwrap, contexts)...)
+        (yε - y0) / ε
+    end
+    return y, ty
+end
+
+function value_and_pushforward(
+    f!::F,
+    y,
+    ::NoPushforwardPrep,
+    backend::AutoSimpleFiniteDiff,
+    x,
+    tx::NTuple{B},
+    contexts::Vararg{Context,C},
+) where {F,B,C}
+    ε = eltype(x)(backend.ε)
+    ty = map(tx) do dx
+        f!(y, x + ε * dx, map(unwrap, contexts)...)
+        yε = copy(y)
+        f!(y, x, map(unwrap, contexts)...)
+        y0 = copy(y)
+        (yε - y0) / ε
+    end
+    f!(y, x, map(unwrap, contexts)...)
+    return y, ty
+end

DifferentiationInterface/test/Misc/FromPrimitive/test.jl → DifferentiationInterface/test/Misc/SimpleFiniteDiff/test.jl

@@ -1,38 +1,31 @@
 using DifferentiationInterface, DifferentiationInterfaceTest
-using DifferentiationInterface: AutoForwardFromPrimitive, AutoReverseFromPrimitive
+using DifferentiationInterface: AutoSimpleFiniteDiff, AutoReverseFromPrimitive
 using DifferentiationInterfaceTest
-using ForwardDiff: ForwardDiff
 using Test
 
 LOGGING = get(ENV, "CI", "false") == "false"
 
 backends = [ #
-    AutoForwardFromPrimitive(AutoForwardDiff(; chunksize=5)),
-    AutoReverseFromPrimitive(AutoForwardDiff(; chunksize=4)),
+    AutoSimpleFiniteDiff(; chunksize=5),
+    AutoReverseFromPrimitive(AutoSimpleFiniteDiff(; chunksize=4)),
 ]
 
 second_order_backends = [ #
     SecondOrder(
-        AutoForwardFromPrimitive(AutoForwardDiff(; chunksize=5)),
-        AutoReverseFromPrimitive(AutoForwardDiff(; chunksize=4)),
+        AutoSimpleFiniteDiff(; chunksize=5),
+        AutoReverseFromPrimitive(AutoSimpleFiniteDiff(; chunksize=4)),
     ),
     SecondOrder(
-        AutoReverseFromPrimitive(AutoForwardDiff(; chunksize=5)),
-        AutoForwardFromPrimitive(AutoForwardDiff(; chunksize=4)),
+        AutoReverseFromPrimitive(AutoSimpleFiniteDiff(; chunksize=5)),
+        AutoSimpleFiniteDiff(; chunksize=4),
     ),
 ]
 
 adaptive_backends = [ #
-    AutoForwardFromPrimitive(AutoForwardDiff()),
-    AutoReverseFromPrimitive(AutoForwardDiff()),
-    SecondOrder(
-        AutoForwardFromPrimitive(AutoForwardDiff()),
-        AutoReverseFromPrimitive(AutoForwardDiff()),
-    ),
-    SecondOrder(
-        AutoReverseFromPrimitive(AutoForwardDiff()),
-        AutoForwardFromPrimitive(AutoForwardDiff()),
-    ),
+    AutoSimpleFiniteDiff(),
+    AutoReverseFromPrimitive(AutoSimpleFiniteDiff()),
+    SecondOrder(AutoSimpleFiniteDiff(), AutoReverseFromPrimitive(AutoSimpleFiniteDiff())),
+    SecondOrder(AutoReverseFromPrimitive(AutoSimpleFiniteDiff()), AutoSimpleFiniteDiff()),
 ]
 
 for backend in vcat(backends, second_order_backends)

DifferentiationInterfaceTest/src/scenarios/scenario.jl

@@ -147,3 +147,7 @@ function adapt_batchsize(backend::AbstractADType, scen::Scenario)
         return backend
     end
 end
+
+function no_matrices(scens::AbstractVector{<:Scenario})
+    return filter(s -> !isa(s.x, AbstractMatrix) && !isa(s.y, AbstractMatrix), scens)
+end

DifferentiationInterfaceTest/test/weird.jl

@@ -39,7 +39,10 @@ static_scenarios(;
 ## Weird arrays
 
 test_differentiation(
-    AutoForwardDiff(), static_scenarios(); benchmark=:full, logging=LOGGING
+    AutoForwardDiff(),
+    DIT.no_matrices(static_scenarios());
+    benchmark=:prepared,
+    logging=LOGGING,
 )
 
 test_differentiation(AutoForwardDiff(), component_scenarios(); logging=LOGGING)

DifferentiationInterfaceTest/test/zero_backends.jl

@@ -2,7 +2,7 @@ using ADTypes
 using DifferentiationInterface
 using DifferentiationInterface: AutoZeroForward, AutoZeroReverse
 using DifferentiationInterfaceTest
-using DifferentiationInterfaceTest: allocfree_scenarios
+using DifferentiationInterfaceTest: allocfree_scenarios, no_matrices
 using Test
 
 LOGGING = get(ENV, "CI", "false") == "false"
@@ -29,17 +29,16 @@ test_differentiation(
 
 data0 = benchmark_differentiation(
     AutoZeroForward(),
-    default_scenarios(; include_batchified=false, include_constantified=true);
+    no_matrices(default_scenarios(; include_batchified=false, include_constantified=true));
     logging=LOGGING,
-    benchmark_seconds=0.1,
 );
 
 data1 = benchmark_differentiation(
     AutoZeroForward(),
-    default_scenarios(; include_batchified=false);
+    no_matrices(default_scenarios(; include_batchified=false));
     benchmark=:full,
     logging=LOGGING,
-    benchmark_seconds=0.2,
+    benchmark_seconds=0.05,
     benchmark_aggregation=maximum,
 );
 
@@ -48,7 +47,7 @@ ADTypes.mode(::FakeBackend) = ADTypes.ForwardMode()
 
 data2 = benchmark_differentiation(
     FakeBackend(),
-    default_scenarios(; include_batchified=false);
+    no_matrices(default_scenarios(; include_batchified=false));
     logging=false,
     benchmark_test=false,
 );