test: make CI faster (#655)

* Make CI faster

* Smaller epsilon

* Use central finite diff

* Import

* Fix batch size

* Fix batch size

Author: gdalle

.github/workflows/Test.yml

@@ -31,11 +31,11 @@ jobs:
           - "1.10"
           - "1"
         group:
-          - Misc/Internals
-          - Misc/DifferentiateWith
-          - Misc/FromPrimitive
-          - Misc/SparsityDetector
-          - Misc/ZeroBackends
+          - Core/Internals
+          - Back/DifferentiateWith
+          - Core/SimpleFiniteDiff
+          - Back/SparsityDetector
+          - Core/ZeroBackends
           - Back/ChainRules
           # - Back/Diffractor
           - Back/Enzyme

DifferentiationInterface/Project.toml

@@ -1,7 +1,7 @@
 name = "DifferentiationInterface"
 uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.6.24"
+version = "0.6.25"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/utils.jl

@@ -1,9 +1,7 @@
 # until https://github.com/EnzymeAD/Enzyme.jl/pull/1545 is merged
 function DI.BatchSizeSettings(::AutoEnzyme, N::Integer)
     B = DI.reasonable_batchsize(N, 16)
-    singlebatch = B == N
-    aligned = N % B == 0
-    return DI.BatchSizeSettings{B,singlebatch,aligned}(N)
+    return DI.BatchSizeSettings{B}(N)
 end
 
 to_val(::DI.BatchSizeSettings{B}) where {B} = Val(B)

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/utils.jl

@@ -1,18 +1,10 @@
 function DI.BatchSizeSettings(::AutoForwardDiff{nothing}, N::Integer)
-    B = ForwardDiff.pickchunksize(N)
-    singlebatch = B == N
-    aligned = N % B == 0
-    return DI.BatchSizeSettings{B,singlebatch,aligned}(N)
+    chunksize = ForwardDiff.pickchunksize(N)
+    return DI.BatchSizeSettings{chunksize}(N)
 end
 
 function DI.BatchSizeSettings(::AutoForwardDiff{chunksize}, N::Integer) where {chunksize}
-    if chunksize > N
-        throw(ArgumentError("Fixed chunksize $chunksize larger than input size $N"))
-    end
-    B = chunksize
-    singlebatch = B == N
-    aligned = N % B == 0
-    return DI.BatchSizeSettings{B,singlebatch,aligned}(N)
+    return DI.BatchSizeSettings{chunksize}(N)
 end
 
 function DI.threshold_batchsize(

DifferentiationInterface/ext/DifferentiationInterfaceStaticArraysExt/DifferentiationInterfaceStaticArraysExt.jl

@@ -14,6 +14,10 @@ end
 
 DI.ismutable_array(::Type{<:SArray}) = false
 
+function DI.BatchSizeSettings(::DI.AutoSimpleFiniteDiff{nothing}, x::StaticArray)
+    return DI.BatchSizeSettings{length(x),true,true}(length(x))
+end
+
 function DI.BatchSizeSettings(::AutoForwardDiff{nothing}, x::StaticArray)
     return DI.BatchSizeSettings{length(x),true,true}(length(x))
 end
@@ -22,4 +26,16 @@ function DI.BatchSizeSettings(::AutoEnzyme, x::StaticArray)
     return DI.BatchSizeSettings{length(x),true,true}(length(x))
 end
 
+function DI.BatchSizeSettings(
+    ::DI.AutoSimpleFiniteDiff{chunksize}, x::StaticArray
+) where {chunksize}
+    return DI.BatchSizeSettings{chunksize}(Val(length(x)))
+end
+
+function DI.BatchSizeSettings(
+    ::AutoForwardDiff{chunksize}, x::StaticArray
+) where {chunksize}
+    return DI.BatchSizeSettings{chunksize}(Val(length(x)))
+end
+
 end

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,87 @@
+"""
+    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 = reasonable_batchsize(N, 12)
+    return BatchSizeSettings{B}(N)
+end
+
+function BatchSizeSettings(::AutoSimpleFiniteDiff{chunksize}, N::Integer) where {chunksize}
+    return BatchSizeSettings{chunksize}(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
+        y1 = f(x + ε * dx, map(unwrap, contexts)...)
+        y0 = f(x - ε * dx, map(unwrap, contexts)...)
+        (y1 - y0) / 2ε
+    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)...)
+        y1 = copy(y)
+        f!(y, x - ε * dx, map(unwrap, contexts)...)
+        y0 = copy(y)
+        (y1 - y0) / 2ε
+    end
+    f!(y, x, map(unwrap, contexts)...)
+    return y, ty
+end

DifferentiationInterface/src/utils/batchsize.jl

@@ -6,7 +6,7 @@ Configuration for the batch size deduced from a backend and a sample array of le
 # Type parameters
 
 - `B::Int`: batch size
-- `singlebatch::Bool`: whether `B > N`
+- `singlebatch::Bool`: whether `B == N` (`B > N` is not allowed)
 - `aligned::Bool`: whether `N % B == 0`
 
 # Fields
@@ -22,11 +22,25 @@ struct BatchSizeSettings{B,singlebatch,aligned}
 end
 
 function BatchSizeSettings{B,singlebatch,aligned}(N::Integer) where {B,singlebatch,aligned}
+    B > N && throw(ArgumentError("Batch size $B larger than input size $N"))
     A = div(N, B, RoundUp)
     B_last = N % B
     return BatchSizeSettings{B,singlebatch,aligned}(N, A, B_last)
 end
 
+function BatchSizeSettings{B}(::Val{N}) where {B,N}
+    singlebatch = B == N
+    aligned = N % B == 0
+    return BatchSizeSettings{B,singlebatch,aligned}(N)
+end
+
+function BatchSizeSettings{B}(N::Integer) where {B}
+    # type-unstable
+    singlebatch = B == N
+    aligned = N % B == 0
+    return BatchSizeSettings{B,singlebatch,aligned}(N)
+end
+
 function BatchSizeSettings(::AbstractADType, N::Integer)
     B = 1
     singlebatch = false