Set default batch size to 1 (#340)

* Default batch size to 1

* Add test

* Remove ambiguity

* Default Enzyme batch size to 8

* Actuallly 16

* Avoid ambiguity

Author: gdalle

DifferentiationInterface/Project.toml

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

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/DifferentiationInterfaceEnzymeExt.jl

@@ -57,6 +57,9 @@ reverse_mode(::AnyAutoEnzyme{Nothing}) = Reverse
 
 DI.check_available(::AutoEnzyme) = true
 
+# until https://github.com/EnzymeAD/Enzyme.jl/pull/1545 is merged
+DI.pick_batchsize(::AnyAutoEnzyme, dimension::Integer) = min(dimension, 16)
+
 # Enzyme's `Duplicated(x, dx)` expects both arguments to be of the same type
 function DI.basis(::AutoEnzyme, a::AbstractArray{T}, i::CartesianIndex) where {T}
     b = zero(a)

DifferentiationInterface/src/misc/from_primitive.jl

@@ -3,6 +3,10 @@ abstract type FromPrimitive <: AbstractADType end
 check_available(fromprim::FromPrimitive) = check_available(fromprim.backend)
 twoarg_support(fromprim::FromPrimitive) = twoarg_support(fromprim.backend)
 
+function pick_batchsize(fromprim::FromPrimitive, dimension::Integer)
+    return pick_batchsize(fromprim.backend, dimension)
+end
+
 ## Forward
 
 struct AutoForwardFromPrimitive{B} <: FromPrimitive

DifferentiationInterface/src/utils/batch.jl

@@ -2,10 +2,10 @@
     pick_batchsize(backend::AbstractADType, dimension::Integer)
 
 Pick a reasonable batch size for batched derivative evaluation with a given total `dimension`.
+
+Returns `1` for backends which have not overloaded it.
 """
-function pick_batchsize(::AbstractADType, dimension::Integer)
-    return min(dimension, 8)
-end
+pick_batchsize(::AbstractADType, dimension::Integer) = 1
 
 """
     Batch{B,T}

DifferentiationInterface/test/Internals/autosparse.jl

@@ -0,0 +1,36 @@
+using ADTypes
+using DifferentiationInterface
+import DifferentiationInterface as DI
+using Test
+
+@testset "SecondOrder" begin
+    backend = SecondOrder(AutoForwardDiff(), AutoZygote())
+    @test ADTypes.mode(backend) isa ADTypes.ForwardMode
+    @test DifferentiationInterface.outer(backend) isa AutoForwardDiff
+    @test DifferentiationInterface.inner(backend) isa AutoZygote
+end
+
+@testset "Sparse" begin
+    for backend in [AutoForwardDiff(), AutoZygote()]
+        sparse_backend = AutoSparse(backend)
+        @test ADTypes.mode(sparse_backend) == ADTypes.mode(backend)
+        @test check_available(sparse_backend) == check_available(backend)
+        @test DI.twoarg_support(sparse_backend) == DI.twoarg_support(backend)
+        @test DI.pushforward_performance(sparse_backend) ==
+            DI.pushforward_performance(backend)
+        @test DI.pullback_performance(sparse_backend) == DI.pullback_performance(backend)
+    end
+
+    for backend in [
+        SecondOrder(AutoForwardDiff(), AutoZygote()),
+        SecondOrder(AutoZygote(), AutoForwardDiff()),
+    ]
+        sparse_backend = AutoSparse(backend)
+        @test ADTypes.mode(sparse_backend) == ADTypes.mode(backend)
+        @test DI.hvp_mode(sparse_backend) == DI.hvp_mode(backend)
+    end
+end
+
+@testset "Batch size" begin
+    @test DI.pick_batchsize(AutoZygote(), 2) == 1
+end

DifferentiationInterface/test/Internals/backends.jl

@@ -13,6 +13,7 @@ for backend in vcat(fromprimitive_backends)
     @test check_available(backend)
     @test check_twoarg(backend)
     @test check_hessian(backend)
+    @test DifferentiationInterface.pick_batchsize(backend, 100) == 5
 end
 
 ## Dense backends