Implement pick_batchsize

src/utils.jl

@@ -1,3 +1,19 @@
+"""
+    pick_batchsize(totalsize, mode, ftype, return_activity, argtypes...)
+
+Return a reasonable batch size for batched differentiation.
+
+!!! warning
+    This function is experimental, and not part of the public API.
+"""
+function pick_batchsize(totalsize::Integer,
+                        mode::Mode,
+                        ftype::Type,
+                        return_activity, ::Type{<:Annotation},
+                        argtypes::Vararg{Type{<:Annotation},Nargs}) where {Nargs}
+    return min(totalsize, 16)
+end
+
 """
     unsafe_to_pointer
 
@@ -260,8 +276,3 @@ end
 end
 
 export codegen_world_age
-
-
-
-
-

test/runtests.jl

@@ -84,6 +84,7 @@ Aqua.test_all(Enzyme, unbound_args=false, piracies=false, deps_compat=false)
 
 include("abi.jl")
 include("typetree.jl")
+include("utils.jl")
 
 @static if Enzyme.EnzymeRules.issupported()
     include("rules.jl")
@@ -3516,4 +3517,3 @@ end
 
 
 end
-

test/utils.jl

@@ -0,0 +1,10 @@
+using Enzyme: Active, Duplicated, Forward, pick_batchsize
+using Test
+
+mode = Forward
+ftype = typeof(sum)
+argtypes = typeof.((Duplicated(ones(1), zeros(1)),))
+@test pick_batchsize(1, mode, ftype, Active, argtypes...) == 1
+
+argtypes = typeof.((Duplicated(ones(100), zeros(100)),))
+@test pick_batchsize(100, mode, ftype, Active, argtypes...) == 16