test: make CI faster (#655)

* Make CI faster

* Smaller epsilon

* Use central finite diff

* Import

* Fix batch size

* Fix batch size

.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

DifferentiationInterface/test/Back/ForwardDiff/test.jl

@@ -4,9 +4,10 @@ Pkg.add("ForwardDiff")
 using ADTypes: ADTypes
 using ComponentArrays: ComponentArrays
 using DifferentiationInterface, DifferentiationInterfaceTest
+import DifferentiationInterface as DI
 import DifferentiationInterfaceTest as DIT
 using ForwardDiff: ForwardDiff
-using StaticArrays: StaticArrays
+using StaticArrays: StaticArrays, @SVector
 using Test
 
 using ExplicitImports
@@ -75,7 +76,18 @@ test_differentiation(
 
 test_differentiation(AutoForwardDiff(), static_scenarios(); logging=LOGGING)
 
-@testset verbose = true "No allocations on StaticArrays" begin
+@testset verbose = true "StaticArrays" begin
+    @testset "Batch size" begin
+        @test DI.pick_batchsize(AutoForwardDiff(), rand(7)) isa DI.BatchSizeSettings{7}
+        @test DI.pick_batchsize(AutoForwardDiff(; chunksize=5), rand(7)) isa
+            DI.BatchSizeSettings{5}
+        @test (@inferred DI.pick_batchsize(AutoForwardDiff(), @SVector(rand(7)))) isa
+            DI.BatchSizeSettings{7}
+        @test (@inferred DI.pick_batchsize(
+            AutoForwardDiff(; chunksize=5), @SVector(rand(7))
+        )) isa DI.BatchSizeSettings{5}
+    end
+
     filtered_static_scenarios = filter(static_scenarios(; include_batchified=false)) do scen
         DIT.function_place(scen) == :out && DIT.operator_place(scen) == :out
     end

DifferentiationInterface/test/Core/Internals/batchsize.jl

@@ -0,0 +1,101 @@
+using ADTypes
+using DifferentiationInterface
+using DifferentiationInterface:
+    AutoSimpleFiniteDiff,
+    BatchSizeSettings,
+    pick_batchsize,
+    reasonable_batchsize,
+    threshold_batchsize
+import DifferentiationInterface as DI
+using StaticArrays
+using Test
+
+BSS = BatchSizeSettings
+
+@testset "Default" begin
+    @test (@inferred pick_batchsize(AutoZygote(), zeros(2))) isa BSS{1,false,true}
+    @test (@inferred pick_batchsize(AutoZygote(), zeros(100))) isa BSS{1,false,true}
+    @test_throws ArgumentError pick_batchsize(AutoSparse(AutoZygote()), zeros(2))
+    @test_throws ArgumentError pick_batchsize(
+        SecondOrder(AutoZygote(), AutoZygote()), zeros(2)
+    )
+    @test_throws ArgumentError pick_batchsize(
+        MixedMode(AutoSimpleFiniteDiff(), AutoZygote()), zeros(2)
+    )
+end
+
+@testset "SimpleFiniteDiff (adaptive)" begin
+    @test (pick_batchsize(AutoSimpleFiniteDiff(), zeros(2))) isa BSS{2,true,true}
+    @test (pick_batchsize(AutoSimpleFiniteDiff(), zeros(6))) isa BSS{6,true,true}
+    @test (pick_batchsize(AutoSimpleFiniteDiff(), zeros(12))) isa BSS{12,true,true}
+    @test (pick_batchsize(AutoSimpleFiniteDiff(), zeros(24))) isa BSS{12,false,true}
+    @test (pick_batchsize(AutoSimpleFiniteDiff(), zeros(100))) isa BSS{12,false,false}
+    @test (@inferred pick_batchsize(AutoSimpleFiniteDiff(), @SVector(zeros(2)))) isa
+        BSS{2,true,true}
+    @test (@inferred pick_batchsize(AutoSimpleFiniteDiff(), @SVector(zeros(6)))) isa
+        BSS{6,true,true}
+    @test (@inferred pick_batchsize(AutoSimpleFiniteDiff(), @SVector(zeros(100)))) isa
+        BSS{100,true,true}
+end
+
+@testset "SimpleFiniteDiff (fixed)" begin
+    @test_throws ArgumentError pick_batchsize(AutoSimpleFiniteDiff(; chunksize=4), zeros(2))
+    @test_throws ArgumentError pick_batchsize(
+        AutoSimpleFiniteDiff(; chunksize=4), @SVector(zeros(2))
+    )
+    @test pick_batchsize(AutoSimpleFiniteDiff(; chunksize=4), zeros(6)) isa BSS{4}
+    @test pick_batchsize(AutoSimpleFiniteDiff(; chunksize=4), zeros(100)) isa BSS{4}
+    BSS{4,true,true}
+    @test pick_batchsize(AutoSimpleFiniteDiff(; chunksize=4), zeros(99)) isa BSS{4}
+    BSS{4,true,false}
+    @test (@inferred pick_batchsize(
+        AutoSimpleFiniteDiff(; chunksize=4), @SVector(zeros(6))
+    )) isa BSS{4,false,false}
+    @test (@inferred pick_batchsize(
+        AutoSimpleFiniteDiff(; chunksize=4), @SVector(zeros(100))
+    )) isa BSS{4,false,true}
+end
+
+@testset "Thresholding" begin
+    @test threshold_batchsize(AutoSimpleFiniteDiff(), 2) isa AutoSimpleFiniteDiff{nothing}
+    @test threshold_batchsize(AutoSimpleFiniteDiff(; chunksize=4), 2) isa
+        AutoSimpleFiniteDiff{2}
+    @test threshold_batchsize(AutoSimpleFiniteDiff(; chunksize=4), 6) isa
+        AutoSimpleFiniteDiff{4}
+    @test threshold_batchsize(AutoSparse(AutoSimpleFiniteDiff(; chunksize=4)), 2) isa
+        AutoSparse{<:AutoSimpleFiniteDiff{2}}
+    @test threshold_batchsize(
+        SecondOrder(
+            AutoSimpleFiniteDiff(; chunksize=4), AutoSimpleFiniteDiff(; chunksize=3)
+        ),
+        6,
+    ) isa SecondOrder{<:AutoSimpleFiniteDiff{4},<:AutoSimpleFiniteDiff{3}}
+    @test threshold_batchsize(
+        SecondOrder(
+            AutoSimpleFiniteDiff(; chunksize=4), AutoSimpleFiniteDiff(; chunksize=3)
+        ),
+        2,
+    ) isa SecondOrder{<:AutoSimpleFiniteDiff{2},<:AutoSimpleFiniteDiff{2}}
+    @test threshold_batchsize(
+        SecondOrder(
+            AutoSimpleFiniteDiff(; chunksize=1), AutoSimpleFiniteDiff(; chunksize=3)
+        ),
+        2,
+    ) isa SecondOrder{<:AutoSimpleFiniteDiff{1},<:AutoSimpleFiniteDiff{2}}
+    @test threshold_batchsize(
+        SecondOrder(
+            AutoSimpleFiniteDiff(; chunksize=4), AutoSimpleFiniteDiff(; chunksize=1)
+        ),
+        2,
+    ) isa SecondOrder{<:AutoSimpleFiniteDiff{2},<:AutoSimpleFiniteDiff{1}}
+end
+
+@testset "Reasonable" begin
+    for Bmax in 1:5
+        @test all(<=(Bmax), reasonable_batchsize.(1:10, Bmax))
+        @test issorted(div.(1:10, reasonable_batchsize.(1:10, Bmax), RoundUp))
+        if Bmax > 2
+            @test reasonable_batchsize(Bmax + 1, Bmax) < Bmax
+        end
+    end
+end