From 8335c1cfaaeab1c1e9135a595dfd6fbb0c734fbf Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Wed, 8 Oct 2025 13:45:41 +0200
Subject: [PATCH 01/10] Add chunk size selection strategies

---
 Project.toml                     |  2 +-
 lib/EnzymeCore/Project.toml      |  2 +-
 lib/EnzymeCore/src/EnzymeCore.jl | 40 ++++++++++++++++++++++++++++++++
 lib/EnzymeCore/test/chunk.jl     | 12 ++++++++++
 lib/EnzymeCore/test/runtests.jl  | 17 +++++++-------
 src/Enzyme.jl                    |  2 ++
 src/sugar.jl                     |  8 +++++--
 test/sugar.jl                    |  9 +++++++
 8 files changed, 80 insertions(+), 12 deletions(-)
 create mode 100644 lib/EnzymeCore/test/chunk.jl

diff --git a/Project.toml b/Project.toml
index de0129a983..ee60644e73 100644
--- a/Project.toml
+++ b/Project.toml
@@ -44,7 +44,7 @@ BFloat16s = "0.2, 0.3, 0.4, 0.5"
 CEnum = "0.4, 0.5"
 ChainRulesCore = "1"
 DynamicPPL = "0.35, 0.36, 0.37"
-EnzymeCore = "0.8.14"
+EnzymeCore = "0.8.15"
 Enzyme_jll = "0.0.203"
 GPUArraysCore = "0.1.6, 0.2"
 GPUCompiler = "1.6.2"
diff --git a/lib/EnzymeCore/Project.toml b/lib/EnzymeCore/Project.toml
index 7aa6f44aa9..a599346b09 100644
--- a/lib/EnzymeCore/Project.toml
+++ b/lib/EnzymeCore/Project.toml
@@ -1,7 +1,7 @@
 name = "EnzymeCore"
 uuid = "f151be2c-9106-41f4-ab19-57ee4f262869"
 authors = ["William Moses <wmoses@mit.edu>", "Valentin Churavy <vchuravy@mit.edu>"]
-version = "0.8.14"
+version = "0.8.15"
 
 [compat]
 Adapt = "3, 4"
diff --git a/lib/EnzymeCore/src/EnzymeCore.jl b/lib/EnzymeCore/src/EnzymeCore.jl
index 52d211dafc..7cb81d70d9 100644
--- a/lib/EnzymeCore/src/EnzymeCore.jl
+++ b/lib/EnzymeCore/src/EnzymeCore.jl
@@ -8,6 +8,7 @@ export DefaultABI, FFIABI, InlineABI, NonGenABI
 export BatchDuplicatedFunc
 export within_autodiff, ignore_derivatives
 export needs_primal
+export ChunkStrategy, OneChunk, AutoChunk, pick_chunksize
 
 function batch_size end
 
@@ -797,4 +798,43 @@ end
 
 Combined(mode::ReverseMode) = mode
 
+"""
+    ChunkStrategy
+
+Abstract type gathering strategies for chunk size selection.
+
+# See also
+
+- [`OneChunk`](@ref)
+- [`AutoChunk`](@ref)
+"""
+abstract type ChunkStrategy end
+
+"""
+    OneChunk()
+
+Select chunk size so that the corresponding array is processed in a single chunk.
+"""
+struct OneChunk <: ChunkStrategy end
+
+"""
+    AutoChunk()
+
+Select chunk size automatically based on internal Enzyme-specific heuristics.
+"""
+struct AutoChunk <: ChunkStrategy end
+
+const DEFAULT_CHUNK_SIZE = 16
+
+"""
+    pick_chunksize(s::ChunkStrategy, a::AbstractArray)
+
+Return the chunk size chosen by strategy `s` based on the dimension of array `a`.
+
+- In forward-mode gradients and Jacobians, `a` would be the input array.
+- In reverse-mode Jacobians, `a` would be the output array.
+"""
+pick_chunksize(::OneChunk, a::AbstractArray) = length(a)
+pick_chunksize(::AutoChunk, a::AbstractArray) = min(DEFAULT_CHUNK_SIZE, length(a))
+
 end # module EnzymeCore
diff --git a/lib/EnzymeCore/test/chunk.jl b/lib/EnzymeCore/test/chunk.jl
new file mode 100644
index 0000000000..edc66029bf
--- /dev/null
+++ b/lib/EnzymeCore/test/chunk.jl
@@ -0,0 +1,12 @@
+using Test
+using EnzymeCore
+
+@testset "OneChunk" begin
+    @test pick_chunksize(OneChunk(), ones(10)) == 10
+    @test pick_chunksize(OneChunk(), ones(100)) == 100
+end
+
+@testset "AutoChunk" begin
+    @test pick_chunksize(AutoChunk(), ones(10)) == 10
+    @test pick_chunksize(AutoChunk(), ones(100)) == 16
+end
diff --git a/lib/EnzymeCore/test/runtests.jl b/lib/EnzymeCore/test/runtests.jl
index c32747e0f0..bf2bcd2143 100644
--- a/lib/EnzymeCore/test/runtests.jl
+++ b/lib/EnzymeCore/test/runtests.jl
@@ -36,12 +36,13 @@ using EnzymeCore
     @testset "Mode modification" begin
         include("mode_modification.jl")
     end
-end
-
-@testset "within_autodiff" begin
-    @test !EnzymeCore.within_autodiff()
-end
-
-@testset "ignore_derivatives" begin
-    @test EnzymeCore.ignore_derivatives(3) == 3
+    @testset "Chunk strategy" begin
+        include("chunk.jl")
+    end
+    @testset "within_autodiff" begin
+        @test !EnzymeCore.within_autodiff()
+    end
+    @testset "ignore_derivatives" begin
+        @test EnzymeCore.ignore_derivatives(3) == 3
+    end
 end
diff --git a/src/Enzyme.jl b/src/Enzyme.jl
index 595bbbbe90..5ad1f8c451 100644
--- a/src/Enzyme.jl
+++ b/src/Enzyme.jl
@@ -111,6 +111,8 @@ export autodiff,
     make_zero!,
     remake_zero!
 
+import EnzymeCore: ChunkStrategy, pick_chunksize
+
 export jacobian, gradient, gradient!, hvp, hvp!, hvp_and_gradient!
 export batch_size, onehot, chunkedonehot
 
diff --git a/src/sugar.jl b/src/sugar.jl
index edec81e052..0aee9d5926 100644
--- a/src/sugar.jl
+++ b/src/sugar.jl
@@ -428,6 +428,10 @@ end
     return ((one(x),),)
 end
 
+@inline function chunkedonehot(x, strategy::ChunkStrategy)
+    return chunkedonehot(x, Val(pick_chunksize(strategy, x)))
+end
+
 @inline tupleconcat(x) = x
 @inline tupleconcat(x, y) = (x..., y...)
 @inline tupleconcat(x, y, z...) = (x..., tupleconcat(y, z...)...)
@@ -712,7 +716,7 @@ gradient(Forward, mul, [2.0, 3.0], Const([2.7, 3.1]))
                 push!(subderivatives, :(values($resp[1])))
             end
             :(($(subderivatives...),))
-        else
+        else  # TODO: handle OneChunk and MaxChunk
             subderivatives = Union{Symbol,Expr}[]
             for an in 1:argnum
                 dargs = Union{Symbol,Expr}[]
@@ -914,7 +918,7 @@ end
 
     chunksize = if chunk <: Val
         chunk.parameters[1]
-    else
+    else  # TODO: handle OneChunk and MaxChunk
         1
     end
     num = ((n_out_val + chunksize - 1) ÷ chunksize)
diff --git a/test/sugar.jl b/test/sugar.jl
index 3e133248ff..9b8df853ad 100644
--- a/test/sugar.jl
+++ b/test/sugar.jl
@@ -666,3 +666,12 @@ end
     # @show J_r_3(u, A, x)
     # @show J_f_3(u, A, x)
 end
+
+@testset "Chunk size strategies" begin  # not passing yet
+    @test_nowarn gradient(Forward, sum, ones(10); chunk=OneChunk())
+    @test_nowarn gradient(Forward, sum, ones(10); chunk=AutoChunk())
+    @test_nowarn jacobian(Forward, copy, ones(10); chunk=OneChunk())
+    @test_nowarn jacobian(Forward, copy, ones(10); chunk=AutoChunk())
+    @test_nowarn jacobian(Reverse, copy, ones(10); chunk=OneChunk())
+    @test_nowarn jacobian(Reverse, copy, ones(10); chunk=AutoChunk())
+end

From 336fd3ceb5959a23ccb6ffac38ae953eab73d3f6 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Wed, 29 Oct 2025 09:59:13 +0100
Subject: [PATCH 02/10] Passing tests except in reverse mode

---
 lib/EnzymeCore/src/EnzymeCore.jl | 21 ++++++++-------
 lib/EnzymeCore/test/chunk.jl     | 10 +++----
 src/Enzyme.jl                    |  3 ++-
 src/sugar.jl                     | 19 ++++++++------
 test/sugar.jl                    | 45 +++++++++++++++++++++++++++-----
 5 files changed, 68 insertions(+), 30 deletions(-)

diff --git a/lib/EnzymeCore/src/EnzymeCore.jl b/lib/EnzymeCore/src/EnzymeCore.jl
index 7cb81d70d9..c39c3511f6 100644
--- a/lib/EnzymeCore/src/EnzymeCore.jl
+++ b/lib/EnzymeCore/src/EnzymeCore.jl
@@ -8,7 +8,7 @@ export DefaultABI, FFIABI, InlineABI, NonGenABI
 export BatchDuplicatedFunc
 export within_autodiff, ignore_derivatives
 export needs_primal
-export ChunkStrategy, OneChunk, AutoChunk, pick_chunksize
+export ChunkStrategy, SingleChunk, AutoChunk, pick_chunksize
 
 function batch_size end
 
@@ -805,22 +805,22 @@ Abstract type gathering strategies for chunk size selection.
 
 # See also
 
-- [`OneChunk`](@ref)
+- [`SingleChunk`](@ref)
 - [`AutoChunk`](@ref)
 """
 abstract type ChunkStrategy end
 
 """
-    OneChunk()
+    SingleChunk()
 
-Select chunk size so that the corresponding array is processed in a single chunk.
+Select chunk size equal to the number of elements, so that the corresponding array is processed in a single chunk.
 """
-struct OneChunk <: ChunkStrategy end
+struct SingleChunk <: ChunkStrategy end
 
 """
     AutoChunk()
 
-Select chunk size automatically based on internal Enzyme-specific heuristics.
+Select chunk size automatically based on internal Enzyme-specific heuristics, which are subject to change.
 """
 struct AutoChunk <: ChunkStrategy end
 
@@ -829,12 +829,15 @@ const DEFAULT_CHUNK_SIZE = 16
 """
     pick_chunksize(s::ChunkStrategy, a::AbstractArray)
 
-Return the chunk size chosen by strategy `s` based on the dimension of array `a`.
+Return the chunk size chosen by strategy `s` based on the dimension of array `a`, as a `Val{C}` object.
 
 - In forward-mode gradients and Jacobians, `a` would be the input array.
 - In reverse-mode Jacobians, `a` would be the output array.
+
+!!! warning
+    For `SingleChunk` and `AutoChunk` strategies, this function is type-unstable.
 """
-pick_chunksize(::OneChunk, a::AbstractArray) = length(a)
-pick_chunksize(::AutoChunk, a::AbstractArray) = min(DEFAULT_CHUNK_SIZE, length(a))
+pick_chunksize(::SingleChunk, a::AbstractArray) = Val(length(a))
+pick_chunksize(::AutoChunk, a::AbstractArray) = Val(min(DEFAULT_CHUNK_SIZE, length(a)))
 
 end # module EnzymeCore
diff --git a/lib/EnzymeCore/test/chunk.jl b/lib/EnzymeCore/test/chunk.jl
index edc66029bf..fc2c5efbf2 100644
--- a/lib/EnzymeCore/test/chunk.jl
+++ b/lib/EnzymeCore/test/chunk.jl
@@ -1,12 +1,12 @@
 using Test
 using EnzymeCore
 
-@testset "OneChunk" begin
-    @test pick_chunksize(OneChunk(), ones(10)) == 10
-    @test pick_chunksize(OneChunk(), ones(100)) == 100
+@testset "SingleChunk" begin
+    @test pick_chunksize(SingleChunk(), ones(10)) == Val(10)
+    @test pick_chunksize(SingleChunk(), ones(100)) == Val(100)
 end
 
 @testset "AutoChunk" begin
-    @test pick_chunksize(AutoChunk(), ones(10)) == 10
-    @test pick_chunksize(AutoChunk(), ones(100)) == 16
+    @test pick_chunksize(AutoChunk(), ones(10)) == Val(10)
+    @test pick_chunksize(AutoChunk(), ones(100)) == Val(16)
 end
diff --git a/src/Enzyme.jl b/src/Enzyme.jl
index 5ad1f8c451..dcb950f6c2 100644
--- a/src/Enzyme.jl
+++ b/src/Enzyme.jl
@@ -111,10 +111,11 @@ export autodiff,
     make_zero!,
     remake_zero!
 
-import EnzymeCore: ChunkStrategy, pick_chunksize
+import EnzymeCore: ChunkStrategy, SingleChunk, AutoChunk, pick_chunksize
 
 export jacobian, gradient, gradient!, hvp, hvp!, hvp_and_gradient!
 export batch_size, onehot, chunkedonehot
+export SingleChunk, AutoChunk
 
 using LinearAlgebra
 import SparseArrays
diff --git a/src/sugar.jl b/src/sugar.jl
index 0aee9d5926..8142c0611c 100644
--- a/src/sugar.jl
+++ b/src/sugar.jl
@@ -429,7 +429,7 @@ end
 end
 
 @inline function chunkedonehot(x, strategy::ChunkStrategy)
-    return chunkedonehot(x, Val(pick_chunksize(strategy, x)))
+    return chunkedonehot(x, pick_chunksize(strategy, x))
 end
 
 @inline tupleconcat(x) = x
@@ -506,10 +506,11 @@ end
 @inline specialize_output(output, input) = output
 
 """
-    gradient(::ForwardMode, f, x; shadows=onehot(x), chunk=nothing)
+    gradient(::ForwardMode, f, x, args...; chunk=nothing, shadows=create_shadows(chunk, x, args...))
 
-Compute the gradient of an array-input function `f` using forward mode. The
-optional keyword argument `shadow` is a vector of one-hot vectors of type `x`
+Compute the gradient of an array-input function `f` using forward mode.
+The optional keyword argument `chunk` optionally denotes the chunk size to use: it can be either `nothing`, `Val(C)` for some `C`, `SingleChunk()` or `AutoChunk()`.
+The optional keyword argument `shadow` is a vector of one-hot vectors of type `x`
 which are used to forward-propagate into the return. For performance reasons,
 this should be computed once, outside the call to `gradient`, rather than
 within this call.
@@ -716,7 +717,7 @@ gradient(Forward, mul, [2.0, 3.0], Const([2.7, 3.1]))
                 push!(subderivatives, :(values($resp[1])))
             end
             :(($(subderivatives...),))
-        else  # TODO: handle OneChunk and MaxChunk
+        else
             subderivatives = Union{Symbol,Expr}[]
             for an in 1:argnum
                 dargs = Union{Symbol,Expr}[]
@@ -792,7 +793,7 @@ end
 """
     jacobian(::ForwardMode, args...; kwargs...)
 
-Equivalent to gradient(::ForwardMode, args...; kwargs...)
+Equivalent to `gradient(::ForwardMode, args...; kwargs...)`.
 """
 @inline function jacobian(fm::ForwardMode, args...; kwargs...)
     gradient(fm, args...; kwargs...)
@@ -918,7 +919,9 @@ end
 
     chunksize = if chunk <: Val
         chunk.parameters[1]
-    else  # TODO: handle OneChunk and MaxChunk
+    else
+        # TODO: handle SingleChunk and MaxChunk
+        # this will change the generated function because the chunksize might be determined at runtime
         1
     end
     num = ((n_out_val + chunksize - 1) ÷ chunksize)
@@ -1177,7 +1180,7 @@ end
     jacobian(::ReverseMode, f, x)
 
 Compute the jacobian of a array-output function `f` using (potentially vector)
-reverse mode. The `chunk` argument optionally denotes the chunk size to use and
+reverse mode. The `chunk` argument optionally denotes the chunk size to use (it can be either `nothing` or `Val(C)` for some `C`) and
 `n_outs` optionally denotes the shape of the array returned by `f` (e.g `size(f(x))`).
 
 Example:
diff --git a/test/sugar.jl b/test/sugar.jl
index 9b8df853ad..2b19ead62d 100644
--- a/test/sugar.jl
+++ b/test/sugar.jl
@@ -667,11 +667,42 @@ end
     # @show J_f_3(u, A, x)
 end
 
-@testset "Chunk size strategies" begin  # not passing yet
-    @test_nowarn gradient(Forward, sum, ones(10); chunk=OneChunk())
-    @test_nowarn gradient(Forward, sum, ones(10); chunk=AutoChunk())
-    @test_nowarn jacobian(Forward, copy, ones(10); chunk=OneChunk())
-    @test_nowarn jacobian(Forward, copy, ones(10); chunk=AutoChunk())
-    @test_nowarn jacobian(Reverse, copy, ones(10); chunk=OneChunk())
-    @test_nowarn jacobian(Reverse, copy, ones(10); chunk=AutoChunk())
+fchunk1(x) = sum(sin, x)
+fchunk2(x) = map(sin, x) + map(cos, reverse(x))
+
+@testset "Chunking strategies" begin
+    @testset "ChunkedOneHot" begin
+        @test chunkedonehot(ones(10), SingleChunk()) isa Tuple{NTuple{10}}
+        @test chunkedonehot(ones(30), SingleChunk()) isa Tuple{NTuple{30}}
+        @test chunkedonehot(ones(10), AutoChunk()) isa Tuple{NTuple{10}}
+        @test chunkedonehot(ones(30), AutoChunk()) isa Tuple{NTuple{16}, NTuple{14}}
+        @test chunkedonehot(ones(30), AutoChunk()) isa Tuple{NTuple{16}, NTuple{14}}
+        @test chunkedonehot(ones(40), AutoChunk()) isa Tuple{NTuple{16}, NTuple{16}, NTuple{8}}
+    end
+
+    @testset "Forward gradient" begin
+        for n in (10, 30)
+            x = ones(n)
+            g = gradient(Forward, fchunk1, x)
+            @test g == gradient(Forward, fchunk1, x; chunk = SingleChunk())
+            @test g == gradient(Forward, fchunk1, x; chunk = AutoChunk())
+        end
+    end
+    @testset "Forward Jacobian" begin
+        for n in (10, 30)
+            x = ones(n)
+            J = jacobian(Forward, fchunk2, x)
+            @test J == jacobian(Forward, fchunk2, x; chunk = SingleChunk())
+            @test J == jacobian(Forward, fchunk2, x; chunk = AutoChunk())
+        end
+    end
+    @testset "Reverse Jacobian" begin
+        for n in (10, 30)
+            x = ones(n)
+            J = jacobian(Forward, fchunk2, x)
+            # TODO: fix this
+            @test_broken J == jacobian(Reverse, fchunk2, x; chunk = SingleChunk())
+            @test_broken J == jacobian(Reverse, fchunk2, x; chunk = AutoChunk())
+        end
+    end
 end

From a61b08a9e4385fce3cca5030a0dc23ec33ab95fe Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Wed, 29 Oct 2025 20:12:40 +0100
Subject: [PATCH 03/10] Update src/sugar.jl

---
 src/sugar.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/sugar.jl b/src/sugar.jl
index 8142c0611c..8e39ed6c1a 100644
--- a/src/sugar.jl
+++ b/src/sugar.jl
@@ -1180,7 +1180,7 @@ end
     jacobian(::ReverseMode, f, x)
 
 Compute the jacobian of a array-output function `f` using (potentially vector)
-reverse mode. The `chunk` argument optionally denotes the chunk size to use (it can be either `nothing` or `Val(C)` for some `C`) and
+reverse mode. The `chunk` argument optionally denotes the chunk size to use (it can be either `nothing` or `Val(C)` for some integer `C`) and
 `n_outs` optionally denotes the shape of the array returned by `f` (e.g `size(f(x))`).
 
 Example:

From 486529222b882b1784b9dcf387fa4e204f9efe89 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Wed, 29 Oct 2025 20:12:47 +0100
Subject: [PATCH 04/10] Update src/sugar.jl

---
 src/sugar.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/sugar.jl b/src/sugar.jl
index 8e39ed6c1a..68de970294 100644
--- a/src/sugar.jl
+++ b/src/sugar.jl
@@ -509,7 +509,7 @@ end
     gradient(::ForwardMode, f, x, args...; chunk=nothing, shadows=create_shadows(chunk, x, args...))
 
 Compute the gradient of an array-input function `f` using forward mode.
-The optional keyword argument `chunk` optionally denotes the chunk size to use: it can be either `nothing`, `Val(C)` for some `C`, `SingleChunk()` or `AutoChunk()`.
+The optional keyword argument `chunk` denotes the chunk size to use: it can be either `nothing`, `Val(C)` for some integer `C`, `SingleChunk()` or `AutoChunk()`.
 The optional keyword argument `shadow` is a vector of one-hot vectors of type `x`
 which are used to forward-propagate into the return. For performance reasons,
 this should be computed once, outside the call to `gradient`, rather than

From b4078d13bc89c87604519c3b6e24c019f9b51828 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Thu, 30 Oct 2025 10:04:22 +0100
Subject: [PATCH 05/10] Add FixedChunk

---
 lib/EnzymeCore/src/EnzymeCore.jl | 20 +++++++++++++++++++-
 lib/EnzymeCore/test/chunk.jl     |  7 +++++++
 src/Enzyme.jl                    |  4 ++--
 src/sugar.jl                     | 12 ++++++------
 test/sugar.jl                    |  5 +++++
 5 files changed, 39 insertions(+), 9 deletions(-)

diff --git a/lib/EnzymeCore/src/EnzymeCore.jl b/lib/EnzymeCore/src/EnzymeCore.jl
index c39c3511f6..eb9cf158f6 100644
--- a/lib/EnzymeCore/src/EnzymeCore.jl
+++ b/lib/EnzymeCore/src/EnzymeCore.jl
@@ -8,7 +8,7 @@ export DefaultABI, FFIABI, InlineABI, NonGenABI
 export BatchDuplicatedFunc
 export within_autodiff, ignore_derivatives
 export needs_primal
-export ChunkStrategy, SingleChunk, AutoChunk, pick_chunksize
+export ChunkStrategy, SingleChunk, FixedChunk, AutoChunk, pick_chunksize
 
 function batch_size end
 
@@ -806,6 +806,7 @@ Abstract type gathering strategies for chunk size selection.
 # See also
 
 - [`SingleChunk`](@ref)
+- [`FixedChunk`](@ref)
 - [`AutoChunk`](@ref)
 """
 abstract type ChunkStrategy end
@@ -817,6 +818,16 @@ Select chunk size equal to the number of elements, so that the corresponding arr
 """
 struct SingleChunk <: ChunkStrategy end
 
+"""
+    FixedChunk{C}()
+
+Select chunk size equal to a fixed integer `C`.
+
+!!! warning
+    This chunk strategy will error if the corresponding array has length `< C`.
+"""
+struct FixedChunk{C} <: ChunkStrategy end
+
 """
     AutoChunk()
 
@@ -840,4 +851,11 @@ Return the chunk size chosen by strategy `s` based on the dimension of array `a`
 pick_chunksize(::SingleChunk, a::AbstractArray) = Val(length(a))
 pick_chunksize(::AutoChunk, a::AbstractArray) = Val(min(DEFAULT_CHUNK_SIZE, length(a)))
 
+function pick_chunksize(::FixedChunk{C}, a::AbstractArray) where {C}
+    if length(a) < C
+        error("Chunk size $C is larger than array length $(length(a))")
+    end
+    return Val{C}()
+end
+
 end # module EnzymeCore
diff --git a/lib/EnzymeCore/test/chunk.jl b/lib/EnzymeCore/test/chunk.jl
index fc2c5efbf2..c115199319 100644
--- a/lib/EnzymeCore/test/chunk.jl
+++ b/lib/EnzymeCore/test/chunk.jl
@@ -6,6 +6,13 @@ using EnzymeCore
     @test pick_chunksize(SingleChunk(), ones(100)) == Val(100)
 end
 
+@testset "FixedChunk" begin
+    @test_throws ErrorException pick_chunksize(FixedChunk{3}(), ones(2))
+    @test pick_chunksize(FixedChunk{3}(), ones(10)) == Val(3)
+    @test pick_chunksize(FixedChunk{3}(), ones(100)) == Val(3)
+    @test pick_chunksize(FixedChunk{4}(), ones(100)) == Val(4)
+end
+
 @testset "AutoChunk" begin
     @test pick_chunksize(AutoChunk(), ones(10)) == Val(10)
     @test pick_chunksize(AutoChunk(), ones(100)) == Val(16)
diff --git a/src/Enzyme.jl b/src/Enzyme.jl
index 8ad28411b2..7e85e04313 100644
--- a/src/Enzyme.jl
+++ b/src/Enzyme.jl
@@ -111,11 +111,11 @@ export autodiff,
     make_zero!,
     remake_zero!
 
-import EnzymeCore: ChunkStrategy, SingleChunk, AutoChunk, pick_chunksize
+import EnzymeCore: ChunkStrategy, SingleChunk, FixedChunk, AutoChunk, pick_chunksize
 
 export jacobian, gradient, gradient!, hvp, hvp!, hvp_and_gradient!
 export batch_size, onehot, chunkedonehot
-export SingleChunk, AutoChunk
+export SingleChunk, FixedChunk, AutoChunk
 
 using LinearAlgebra
 import SparseArrays
diff --git a/src/sugar.jl b/src/sugar.jl
index 68de970294..fdee4eb7f0 100644
--- a/src/sugar.jl
+++ b/src/sugar.jl
@@ -509,7 +509,7 @@ end
     gradient(::ForwardMode, f, x, args...; chunk=nothing, shadows=create_shadows(chunk, x, args...))
 
 Compute the gradient of an array-input function `f` using forward mode.
-The optional keyword argument `chunk` denotes the chunk size to use: it can be either `nothing`, `Val(C)` for some integer `C`, `SingleChunk()` or `AutoChunk()`.
+The optional keyword argument `chunk` denotes the chunk size to use: it can be either `nothing`, `Val(C)` / `FixedChunk{C}()` for some integer `C` (both are equivalent), `SingleChunk()` or `AutoChunk()`.
 The optional keyword argument `shadow` is a vector of one-hot vectors of type `x`
 which are used to forward-propagate into the return. For performance reasons,
 this should be computed once, outside the call to `gradient`, rather than
@@ -693,7 +693,7 @@ gradient(Forward, mul, [2.0, 3.0], Const([2.7, 3.1]))
             end
 
             :(values($resp[1]))
-        elseif CS == Val{1}
+        elseif CS == Val{1} || CS == FixedChunk{1}
             subderivatives = Union{Symbol,Expr}[]
             for an in 1:argnum
                 dargs = Union{Symbol,Expr}[]
@@ -917,7 +917,7 @@ end
     MDTys = Union{Expr,Symbol}[]
     MDTysLast = Union{Expr,Symbol}[]
 
-    chunksize = if chunk <: Val
+    chunksize = if chunk <: Val || chunk <: FixedChunk
         chunk.parameters[1]
     else
         # TODO: handle SingleChunk and MaxChunk
@@ -1179,9 +1179,9 @@ end
     jacobian(::ReverseMode, f, x; n_outs=nothing, chunk=nothing)
     jacobian(::ReverseMode, f, x)
 
-Compute the jacobian of a array-output function `f` using (potentially vector)
-reverse mode. The `chunk` argument optionally denotes the chunk size to use (it can be either `nothing` or `Val(C)` for some integer `C`) and
-`n_outs` optionally denotes the shape of the array returned by `f` (e.g `size(f(x))`).
+Compute the jacobian of a array-output function `f` using (potentially vector) reverse mode.
+The optional keyword argument `chunk` denotes the chunk size to use: it can be either `nothing` or `Val(C)` / `FixedChunk{C}()` for some integer `C` (both are equivalent).
+The optional keyword argument `n_outs` denotes the shape of the array returned by `f` (e.g `size(f(x))`).
 
 Example:
 
diff --git a/test/sugar.jl b/test/sugar.jl
index 2b19ead62d..980e9d3d8f 100644
--- a/test/sugar.jl
+++ b/test/sugar.jl
@@ -674,6 +674,8 @@ fchunk2(x) = map(sin, x) + map(cos, reverse(x))
     @testset "ChunkedOneHot" begin
         @test chunkedonehot(ones(10), SingleChunk()) isa Tuple{NTuple{10}}
         @test chunkedonehot(ones(30), SingleChunk()) isa Tuple{NTuple{30}}
+        @test chunkedonehot(ones(10), FixedChunk{4}()) isa Tuple{NTuple{4},NTuple{4},NTuple{2}}
+        @test chunkedonehot(ones(10), FixedChunk{5}()) isa Tuple{NTuple{5},NTuple{5}}
         @test chunkedonehot(ones(10), AutoChunk()) isa Tuple{NTuple{10}}
         @test chunkedonehot(ones(30), AutoChunk()) isa Tuple{NTuple{16}, NTuple{14}}
         @test chunkedonehot(ones(30), AutoChunk()) isa Tuple{NTuple{16}, NTuple{14}}
@@ -684,6 +686,7 @@ fchunk2(x) = map(sin, x) + map(cos, reverse(x))
         for n in (10, 30)
             x = ones(n)
             g = gradient(Forward, fchunk1, x)
+            @test g == gradient(Forward, fchunk1, x; chunk = FixedChunk{3}())
             @test g == gradient(Forward, fchunk1, x; chunk = SingleChunk())
             @test g == gradient(Forward, fchunk1, x; chunk = AutoChunk())
         end
@@ -692,6 +695,7 @@ fchunk2(x) = map(sin, x) + map(cos, reverse(x))
         for n in (10, 30)
             x = ones(n)
             J = jacobian(Forward, fchunk2, x)
+            @test J == jacobian(Forward, fchunk2, x; chunk = FixedChunk{3}())
             @test J == jacobian(Forward, fchunk2, x; chunk = SingleChunk())
             @test J == jacobian(Forward, fchunk2, x; chunk = AutoChunk())
         end
@@ -700,6 +704,7 @@ fchunk2(x) = map(sin, x) + map(cos, reverse(x))
         for n in (10, 30)
             x = ones(n)
             J = jacobian(Forward, fchunk2, x)
+            @test J == jacobian(Reverse, fchunk2, x; chunk = FixedChunk{3}())
             # TODO: fix this
             @test_broken J == jacobian(Reverse, fchunk2, x; chunk = SingleChunk())
             @test_broken J == jacobian(Reverse, fchunk2, x; chunk = AutoChunk())

From 6d4ceaf2ec3a3bae77da923952142702181ac5ae Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Thu, 30 Oct 2025 10:45:46 +0100
Subject: [PATCH 06/10] Add SmallestChunk

---
 lib/EnzymeCore/src/EnzymeCore.jl | 58 +++++++++++++++++++++++---------
 lib/EnzymeCore/test/chunk.jl     | 21 ++++++++++--
 src/Enzyme.jl                    |  4 +--
 src/sugar.jl                     | 16 ++++-----
 test/sugar.jl                    | 45 +++++++++++++------------
 5 files changed, 95 insertions(+), 49 deletions(-)

diff --git a/lib/EnzymeCore/src/EnzymeCore.jl b/lib/EnzymeCore/src/EnzymeCore.jl
index eb9cf158f6..26e328120c 100644
--- a/lib/EnzymeCore/src/EnzymeCore.jl
+++ b/lib/EnzymeCore/src/EnzymeCore.jl
@@ -8,7 +8,7 @@ export DefaultABI, FFIABI, InlineABI, NonGenABI
 export BatchDuplicatedFunc
 export within_autodiff, ignore_derivatives
 export needs_primal
-export ChunkStrategy, SingleChunk, FixedChunk, AutoChunk, pick_chunksize
+export ChunkStrategy, SmallestChunk, LargestChunk, FixedChunk, AutoChunk, pick_chunksize
 
 function batch_size end
 
@@ -803,26 +803,40 @@ Combined(mode::ReverseMode) = mode
 
 Abstract type gathering strategies for chunk size selection.
 
-# See also
+# Subtypes
 
-- [`SingleChunk`](@ref)
+- [`SmallestChunk`](@ref)
+- [`LargestChunkk`](@ref)
 - [`FixedChunk`](@ref)
 - [`AutoChunk`](@ref)
 """
 abstract type ChunkStrategy end
 
 """
-    SingleChunk()
+    SmallestChunk()
+
+Select chunk size equal to 1, so that the corresponding array is processed in as many chunks as it has elements.
+
+!!! tip
+    In the current Enzyme interface, this strategy is equivalent to setting `chunk = nothing`.
+"""
+struct SmallestChunk <: ChunkStrategy end
+
+"""
+    LargestChunk()
 
 Select chunk size equal to the number of elements, so that the corresponding array is processed in a single chunk.
 """
-struct SingleChunk <: ChunkStrategy end
+struct LargestChunk <: ChunkStrategy end
 
 """
     FixedChunk{C}()
 
 Select chunk size equal to a fixed integer `C`.
 
+!!! tip
+    In the current Enzyme interface, this chunk strategy is equivalent to setting `chunk = Val(C)`.
+
 !!! warning
     This chunk strategy will error if the corresponding array has length `< C`.
 """
@@ -838,24 +852,38 @@ struct AutoChunk <: ChunkStrategy end
 const DEFAULT_CHUNK_SIZE = 16
 
 """
+    pick_chunksize(s::ChunkStrategy, n::Integer)
     pick_chunksize(s::ChunkStrategy, a::AbstractArray)
 
-Return the chunk size chosen by strategy `s` based on the dimension of array `a`, as a `Val{C}` object.
-
-- In forward-mode gradients and Jacobians, `a` would be the input array.
+Compute tLargestChunkze chosen by strategy `s` based on the integer `n` or the array `a` (`n` corresponds to the array's length)
+Return a `Val{C}` object.
+LargestChunk
+- In forward-modeLargestChunkand Jacobians, `a` would be the input array.
 - In reverse-mode Jacobians, `a` would be the output array.
 
 !!! warning
-    For `SingleChunk` and `AutoChunk` strategies, this function is type-unstable.
+    For `LargestChunk` and `AutoChunk` strategies, this function is type-unstable.
 """
-pick_chunksize(::SingleChunk, a::AbstractArray) = Val(length(a))
-pick_chunksize(::AutoChunk, a::AbstractArray) = Val(min(DEFAULT_CHUNK_SIZE, length(a)))
+function pick_chunksize end
 
-function pick_chunksize(::FixedChunk{C}, a::AbstractArray) where {C}
-    if length(a) < C
-        error("Chunk size $C is larger than array length $(length(a))")
-    end
+pick_chunksize(::SmallestChunk, a_or_n::Union{Integer,AbstractArray}) = Val(1)
+
+pick_chunksize(::LargestChunk, n::Integer) = Val(n)
+pick_chunksize(::LargestChunk, a::AbstractArray) = Val(length(a))  # allows inference on static arrays
+
+pick_chunksize(::AutoChunk, n::Integer) = Val(min(DEFAULT_CHUNK_SIZE, n))  # TODO: improve
+pick_chunksize(s::AutoChunk, a::AbstractArray) = pick_chunksize(s, length(a))
+
+function pick_chunksize(s::FixedChunk{C}, a_or_n::Union{Integer,AbstractArray}) where {C}
+    check_length(s, a_or_n)
     return Val{C}()
 end
 
+function check_length(::FixedChunk{C}, n::Integer) where {C}
+    if n < C
+        error("Chunk size $C is larger than length $n")
+    end
+end
+check_length(s::FixedChunk{C}, a::AbstractArray) where {C} = check_length(s, length(a))
+
 end # module EnzymeCore
diff --git a/lib/EnzymeCore/test/chunk.jl b/lib/EnzymeCore/test/chunk.jl
index c115199319..a3aed39f3a 100644
--- a/lib/EnzymeCore/test/chunk.jl
+++ b/lib/EnzymeCore/test/chunk.jl
@@ -1,19 +1,34 @@
 using Test
 using EnzymeCore
 
-@testset "SingleChunk" begin
-    @test pick_chunksize(SingleChunk(), ones(10)) == Val(10)
-    @test pick_chunksize(SingleChunk(), ones(100)) == Val(100)
+@testset "SmallestChunk" begin
+    @test pick_chunksize(SmallestChunk(), 10) == Val(1)
+    @test pick_chunksize(SmallestChunk(), ones(10)) == Val(1)
+    @test pick_chunksize(SmallestChunk(), 100) == Val(1)
+    @test pick_chunksize(SmallestChunk(), ones(100)) == Val(1)
+end
+
+@testset "LargestChunk" begin
+    @test pick_chunksize(LargestChunk(), 10) == Val(10)
+    @test pick_chunksize(LargestChunk(), ones(10)) == Val(10)
+    @test pick_chunksize(LargestChunk(), 100) == Val(100)
+    @test pick_chunksize(LargestChunk(), ones(100)) == Val(100)
 end
 
 @testset "FixedChunk" begin
+    @test_throws ErrorException pick_chunksize(FixedChunk{3}(), 2)
     @test_throws ErrorException pick_chunksize(FixedChunk{3}(), ones(2))
+    @test pick_chunksize(FixedChunk{3}(), 10) == Val(3)
     @test pick_chunksize(FixedChunk{3}(), ones(10)) == Val(3)
+    @test pick_chunksize(FixedChunk{3}(), 100) == Val(3)
     @test pick_chunksize(FixedChunk{3}(), ones(100)) == Val(3)
+    @test pick_chunksize(FixedChunk{4}(), 100) == Val(4)
     @test pick_chunksize(FixedChunk{4}(), ones(100)) == Val(4)
 end
 
 @testset "AutoChunk" begin
+    @test pick_chunksize(AutoChunk(), 10) == Val(10)
     @test pick_chunksize(AutoChunk(), ones(10)) == Val(10)
+    @test pick_chunksize(AutoChunk(), 100) == Val(16)
     @test pick_chunksize(AutoChunk(), ones(100)) == Val(16)
 end
diff --git a/src/Enzyme.jl b/src/Enzyme.jl
index 7e85e04313..125ac4c57a 100644
--- a/src/Enzyme.jl
+++ b/src/Enzyme.jl
@@ -111,11 +111,11 @@ export autodiff,
     make_zero!,
     remake_zero!
 
-import EnzymeCore: ChunkStrategy, SingleChunk, FixedChunk, AutoChunk, pick_chunksize
+import EnzymeCore: ChunkStrategy, SmallestChunk, LargestChunk, FixedChunk, AutoChunk, pick_chunksize
 
 export jacobian, gradient, gradient!, hvp, hvp!, hvp_and_gradient!
 export batch_size, onehot, chunkedonehot
-export SingleChunk, FixedChunk, AutoChunk
+export SmallestChunk, LargestChunk, FixedChunk, AutoChunk
 
 using LinearAlgebra
 import SparseArrays
diff --git a/src/sugar.jl b/src/sugar.jl
index fdee4eb7f0..f3c637c9d6 100644
--- a/src/sugar.jl
+++ b/src/sugar.jl
@@ -450,7 +450,7 @@ end
             push!(exprs, :(nothing))
         elseif ty <: AbstractFloat
             push!(exprs, :(nothing))
-        elseif ChunkTy == Nothing || ChunkTy == Val{1}
+        elseif ChunkTy == Nothing || ChunkTy == Val{1} || ChunkTy == SmallestChunk || ChunkTy == FixedChunk{1}
             push!(exprs, :(onehot($arg)))
         else
             push!(exprs, :(chunkedonehot($arg, chunk)))
@@ -509,7 +509,7 @@ end
     gradient(::ForwardMode, f, x, args...; chunk=nothing, shadows=create_shadows(chunk, x, args...))
 
 Compute the gradient of an array-input function `f` using forward mode.
-The optional keyword argument `chunk` denotes the chunk size to use: it can be either `nothing`, `Val(C)` / `FixedChunk{C}()` for some integer `C` (both are equivalent), `SingleChunk()` or `AutoChunk()`.
+The optional keyword argument `chunk` denotes the chunk size to use: it can be either `nothing`, `Val(C)` for some integer `C`, or a subtype of [`EnzymeCore.ChunkStrategy`](@ref EnzymeCore.ChunkStrategy).
 The optional keyword argument `shadow` is a vector of one-hot vectors of type `x`
 which are used to forward-propagate into the return. For performance reasons,
 this should be computed once, outside the call to `gradient`, rather than
@@ -664,7 +664,7 @@ gradient(Forward, mul, [2.0, 3.0], Const([2.7, 3.1]))
             :($resp[1])
         elseif argnum == 0
             vals[i]
-        elseif CS == Nothing
+        elseif CS == Nothing || CS == SmallestChunk
             dargs = Union{Symbol,Expr}[]
             for (j, arg2) in enumerate(syms)
                 if i == j
@@ -919,10 +919,10 @@ end
 
     chunksize = if chunk <: Val || chunk <: FixedChunk
         chunk.parameters[1]
-    else
-        # TODO: handle SingleChunk and MaxChunk
-        # this will change the generated function because the chunksize might be determined at runtime
+    elseif chunk == Nothing || chunk == SmallestChunk
         1
+    else  # chunk isa Union{LargestChunk, MaxChunk} and pick_chunksize returns a Val{C}()
+        typeof(pick_chunksize(chunk(), n_out_val)).parameters[1]
     end
     num = ((n_out_val + chunksize - 1) ÷ chunksize)
 
@@ -947,7 +947,7 @@ end
         else
             push!(exprs, Expr(:(=), mdi, :(Compiler.active_reg_nothrow($xti) == Compiler.ActiveState || Compiler.active_reg_nothrow($xti) == Compiler.MixedState)))
 
-            if chunk == Val{1} || chunk == Nothing
+            if chunk == Val{1} || chunk == Nothing || chunk == SmallestChunk || chunk == FixedChunk{1}
                 push!(MDTys, :($mdi ? MixedDuplicated{$xti} : Duplicated{$xti}))
             else
                 push!(MDTys, :($mdi ? BatchMixedDuplicated{$xti, $chunksize} : BatchDuplicated{$xti, $chunksize}))
@@ -1180,7 +1180,7 @@ end
     jacobian(::ReverseMode, f, x)
 
 Compute the jacobian of a array-output function `f` using (potentially vector) reverse mode.
-The optional keyword argument `chunk` denotes the chunk size to use: it can be either `nothing` or `Val(C)` / `FixedChunk{C}()` for some integer `C` (both are equivalent).
+The optional keyword argument `chunk` denotes the chunk size to use: it can be either `nothing`, `Val(C)` for some integer `C`, or a subtype of [`EnzymeCore.ChunkStrategy`](@ref EnzymeCore.ChunkStrategy).
 The optional keyword argument `n_outs` denotes the shape of the array returned by `f` (e.g `size(f(x))`).
 
 Example:
diff --git a/test/sugar.jl b/test/sugar.jl
index 980e9d3d8f..390e55ed39 100644
--- a/test/sugar.jl
+++ b/test/sugar.jl
@@ -672,8 +672,10 @@ fchunk2(x) = map(sin, x) + map(cos, reverse(x))
 
 @testset "Chunking strategies" begin
     @testset "ChunkedOneHot" begin
-        @test chunkedonehot(ones(10), SingleChunk()) isa Tuple{NTuple{10}}
-        @test chunkedonehot(ones(30), SingleChunk()) isa Tuple{NTuple{30}}
+        @test chunkedonehot(ones(3), SmallestChunk()) isa Tuple{NTuple{1},NTuple{1},NTuple{1}}
+        @test chunkedonehot(ones(30), LargestChunk()) isa Tuple{NTuple{30}}
+        @test chunkedonehot(ones(10), LargestChunk()) isa Tuple{NTuple{10}}
+        @test chunkedonehot(ones(30), LargestChunk()) isa Tuple{NTuple{30}}
         @test chunkedonehot(ones(10), FixedChunk{4}()) isa Tuple{NTuple{4},NTuple{4},NTuple{2}}
         @test chunkedonehot(ones(10), FixedChunk{5}()) isa Tuple{NTuple{5},NTuple{5}}
         @test chunkedonehot(ones(10), AutoChunk()) isa Tuple{NTuple{10}}
@@ -682,32 +684,33 @@ fchunk2(x) = map(sin, x) + map(cos, reverse(x))
         @test chunkedonehot(ones(40), AutoChunk()) isa Tuple{NTuple{16}, NTuple{16}, NTuple{8}}
     end
 
+    strategies = [SmallestChunk(), LargestChunk(), FixedChunk{3}(), AutoChunk()]
+
     @testset "Forward gradient" begin
-        for n in (10, 30)
-            x = ones(n)
-            g = gradient(Forward, fchunk1, x)
-            @test g == gradient(Forward, fchunk1, x; chunk = FixedChunk{3}())
-            @test g == gradient(Forward, fchunk1, x; chunk = SingleChunk())
-            @test g == gradient(Forward, fchunk1, x; chunk = AutoChunk())
+        @testset for chunk in strategies
+            for n in (10, 30)
+                x = ones(n)
+                g = gradient(Forward, fchunk1, x)
+                @test g == gradient(Forward, fchunk1, x; chunk)
+            end
         end
     end
     @testset "Forward Jacobian" begin
-        for n in (10, 30)
-            x = ones(n)
-            J = jacobian(Forward, fchunk2, x)
-            @test J == jacobian(Forward, fchunk2, x; chunk = FixedChunk{3}())
-            @test J == jacobian(Forward, fchunk2, x; chunk = SingleChunk())
-            @test J == jacobian(Forward, fchunk2, x; chunk = AutoChunk())
+        @testset for chunk in strategies
+            for n in (10, 30)
+                x = ones(n)
+                J = jacobian(Forward, fchunk2, x)
+                @test J == jacobian(Forward, fchunk2, x; chunk)
+            end
         end
     end
     @testset "Reverse Jacobian" begin
-        for n in (10, 30)
-            x = ones(n)
-            J = jacobian(Forward, fchunk2, x)
-            @test J == jacobian(Reverse, fchunk2, x; chunk = FixedChunk{3}())
-            # TODO: fix this
-            @test_broken J == jacobian(Reverse, fchunk2, x; chunk = SingleChunk())
-            @test_broken J == jacobian(Reverse, fchunk2, x; chunk = AutoChunk())
+        @testset for chunk in strategies
+            for n in (10, 30)
+                x = ones(n)
+                J = jacobian(Forward, fchunk2, x)
+                @test J == jacobian(Reverse, fchunk2, x; chunk)
+            end
         end
     end
 end

From dc5ff05723dc0ee8d4bc19f68d99f9706c252e3d Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Thu, 30 Oct 2025 20:48:51 +0100
Subject: [PATCH 07/10] Add deprecation warning, let everything flow through
 new strategy objects

---
 src/sugar.jl  | 73 +++++++++++++++++++++++++++++++++------------------
 test/sugar.jl | 24 ++++++++++++-----
 2 files changed, 65 insertions(+), 32 deletions(-)

diff --git a/src/sugar.jl b/src/sugar.jl
index f3c637c9d6..6ab645d422 100644
--- a/src/sugar.jl
+++ b/src/sugar.jl
@@ -415,6 +415,27 @@ gradient!(ReverseWithPrimal, dx, f, [2.0, 3.0])
     end
 end
 
+const ExtendedChunkStrategy = Union{ChunkStrategy, Nothing, Val}
+
+# eats and returns a type because generated functions work on argument types
+get_strategy(chunk::Type{CS}) where {CS<:ChunkStrategy} = chunk
+
+function get_strategy(::Type{Nothing})
+    Base.depwarn(
+        "The `chunk=nothing` configuration will be deprecated in a future release. Please use `chunk=SmallestChunk()` instead.",
+        :get_strategy,
+    )
+    return SmallestChunk
+end
+
+function get_strategy(::Type{Val{C}}) where {C}
+    Base.depwarn(
+        "The `chunk=Val(C)` configuration will be deprecated in a future release. Please use `chunk=FixedChunk{C}()` instead.",
+        :get_strategy,
+    )
+    return FixedChunk{C}
+end
+
 @inline function chunkedonehot(x, ::Val{chunk}) where {chunk}
     sz = length(x)
     num = ((sz + chunk - 1) ÷ chunk)
@@ -436,7 +457,8 @@ end
 @inline tupleconcat(x, y) = (x..., y...)
 @inline tupleconcat(x, y, z...) = (x..., tupleconcat(y, z...)...)
 
-@generated function create_shadows(chunk::ChunkTy, x::X, vargs::Vararg{Any,N}) where {ChunkTy, X, N}
+@generated function create_shadows(chunk::ExtendedChunkStrategy, x::X, vargs::Vararg{Any,N}) where {X, N}
+    chunk_strategy = get_strategy(chunk)
     args =  Union{Symbol,Expr}[:x]
     tys =  Type[X]
     for i in 1:N
@@ -450,7 +472,7 @@ end
             push!(exprs, :(nothing))
         elseif ty <: AbstractFloat
             push!(exprs, :(nothing))
-        elseif ChunkTy == Nothing || ChunkTy == Val{1} || ChunkTy == SmallestChunk || ChunkTy == FixedChunk{1}
+        elseif chunk_strategy == SmallestChunk || chunk_strategy == FixedChunk{1}
             push!(exprs, :(onehot($arg)))
         else
             push!(exprs, :(chunkedonehot($arg, chunk)))
@@ -506,10 +528,10 @@ end
 @inline specialize_output(output, input) = output
 
 """
-    gradient(::ForwardMode, f, x, args...; chunk=nothing, shadows=create_shadows(chunk, x, args...))
+    gradient(::ForwardMode, f, x, args...; chunk=SmallestChunk(), shadows=create_shadows(chunk, x, args...))
 
 Compute the gradient of an array-input function `f` using forward mode.
-The optional keyword argument `chunk` denotes the chunk size to use: it can be either `nothing`, `Val(C)` for some integer `C`, or a subtype of [`EnzymeCore.ChunkStrategy`](@ref EnzymeCore.ChunkStrategy).
+The optional keyword argument `chunk` denotes the chunk size to use: it can be any instance of [`EnzymeCore.ChunkStrategy`](@ref EnzymeCore.ChunkStrategy).
 The optional keyword argument `shadow` is a vector of one-hot vectors of type `x`
 which are used to forward-propagate into the return. For performance reasons,
 this should be computed once, outside the call to `gradient`, rather than
@@ -535,7 +557,7 @@ gradient(ForwardWithPrimal, f, [2.0, 3.0])
 ```
 
 ```jldoctest gradfwd
-gradient(Forward, f, [2.0, 3.0]; chunk=Val(1))
+gradient(Forward, f, [2.0, 3.0]; chunk=FixedChunk{1}())
 
 # output
 
@@ -543,7 +565,7 @@ gradient(Forward, f, [2.0, 3.0]; chunk=Val(1))
 ```
 
 ```jldoctest gradfwd
-gradient(ForwardWithPrimal, f, [2.0, 3.0]; chunk=Val(1))
+gradient(ForwardWithPrimal, f, [2.0, 3.0]; chunk=FixedChunk{1}())
 
 # output
 (derivs = ([3.0, 2.0],), val = 6.0)
@@ -592,9 +614,11 @@ gradient(Forward, mul, [2.0, 3.0], Const([2.7, 3.1]))
     f::F,
     x::ty_0,
     args::Vararg{Any,N};
-    chunk::CS = nothing,
+    chunk::ExtendedChunkStrategy = SmallestChunk(),
     shadows::ST = create_shadows(chunk, x, args...),
-) where {F, ReturnPrimal,ABI,ErrIfFuncWritten,RuntimeActivity,StrongZero,CS,ST, ty_0, N}
+) where {F, ReturnPrimal,ABI,ErrIfFuncWritten,RuntimeActivity,StrongZero,ST, ty_0, N}
+
+    chunk_strategy = get_strategy(chunk)
 
     syms = Union{Symbol,Expr}[:x]
     shads = Union{Symbol,Expr}[:(shadows[1])]
@@ -622,7 +646,7 @@ gradient(Forward, mul, [2.0, 3.0], Const([2.7, 3.1]))
         end
     end
 
-    if CS == Val{0}
+    if chunk_strategy == FixedChunk{0}
         return quote
             Base.@_inline_meta
             throw(ErrorException("Cannot differentiate with a batch size of 0"))
@@ -664,7 +688,7 @@ gradient(Forward, mul, [2.0, 3.0], Const([2.7, 3.1]))
             :($resp[1])
         elseif argnum == 0
             vals[i]
-        elseif CS == Nothing || CS == SmallestChunk
+        elseif chunk_strategy == SmallestChunk
             dargs = Union{Symbol,Expr}[]
             for (j, arg2) in enumerate(syms)
                 if i == j
@@ -693,7 +717,7 @@ gradient(Forward, mul, [2.0, 3.0], Const([2.7, 3.1]))
             end
 
             :(values($resp[1]))
-        elseif CS == Val{1} || CS == FixedChunk{1}
+        elseif chunk_strategy == FixedChunk{1}
             subderivatives = Union{Symbol,Expr}[]
             for an in 1:argnum
                 dargs = Union{Symbol,Expr}[]
@@ -803,10 +827,11 @@ end
     mode::ReverseMode{ReturnPrimal},
     RT::RType,
     n_outs::OutType,
-    chunk::CT,
+    chunk::ExtendedChunkStrategy,
     f::F,
     xs::Vararg{Any, Nargs}
-) where {ReturnPrimal,RType, F,Nargs,OutType,CT}
+) where {ReturnPrimal,RType, F,Nargs,OutType}
+    chunk_strategy = get_strategy(chunk)
     fty = if f <: Enzyme.Annotation
         f.parameters[1]
     else
@@ -896,7 +921,7 @@ end
         end
     end
 
-    if chunk == Val{0}
+    if chunk_strategy == FixedChunk{0}
         return quote
             throw(ErrorException("Cannot differentiate with a batch size of 0"))
         end
@@ -917,13 +942,9 @@ end
     MDTys = Union{Expr,Symbol}[]
     MDTysLast = Union{Expr,Symbol}[]
 
-    chunksize = if chunk <: Val || chunk <: FixedChunk
-        chunk.parameters[1]
-    elseif chunk == Nothing || chunk == SmallestChunk
-        1
-    else  # chunk isa Union{LargestChunk, MaxChunk} and pick_chunksize returns a Val{C}()
-        typeof(pick_chunksize(chunk(), n_out_val)).parameters[1]
-    end
+    chunksize_val = pick_chunksize(chunk_strategy(), n_out_val)
+    chunksize = typeof(chunksize_val).parameters[1]
+
     num = ((n_out_val + chunksize - 1) ÷ chunksize)
 
     last_size = if num * chunksize == n_out_val
@@ -947,7 +968,7 @@ end
         else
             push!(exprs, Expr(:(=), mdi, :(Compiler.active_reg_nothrow($xti) == Compiler.ActiveState || Compiler.active_reg_nothrow($xti) == Compiler.MixedState)))
 
-            if chunk == Val{1} || chunk == Nothing || chunk == SmallestChunk || chunk == FixedChunk{1}
+            if chunk_strategy == SmallestChunk || chunk_strategy == FixedChunk{1}
                 push!(MDTys, :($mdi ? MixedDuplicated{$xti} : Duplicated{$xti}))
             else
                 push!(MDTys, :($mdi ? BatchMixedDuplicated{$xti, $chunksize} : BatchDuplicated{$xti, $chunksize}))
@@ -1176,11 +1197,11 @@ end
 end
 
 """
-    jacobian(::ReverseMode, f, x; n_outs=nothing, chunk=nothing)
+    jacobian(::ReverseMode, f, x; n_outs=nothing, chunk=SmallestChunk())
     jacobian(::ReverseMode, f, x)
 
 Compute the jacobian of a array-output function `f` using (potentially vector) reverse mode.
-The optional keyword argument `chunk` denotes the chunk size to use: it can be either `nothing`, `Val(C)` for some integer `C`, or a subtype of [`EnzymeCore.ChunkStrategy`](@ref EnzymeCore.ChunkStrategy).
+The optional keyword argument `chunk` denotes the chunk size to use: it can be any instance of [`EnzymeCore.ChunkStrategy`](@ref EnzymeCore.ChunkStrategy).
 The optional keyword argument `n_outs` denotes the shape of the array returned by `f` (e.g `size(f(x))`).
 
 Example:
@@ -1234,8 +1255,8 @@ this function will retun an AbstractArray of shape `size(output)` of values of t
     f::F,
     xs::Vararg{Any, Nargs};
     n_outs::OutType = nothing,
-    chunk::CT = nothing,
-) where {F,Nargs, OutType,CT}
+    chunk::ExtendedChunkStrategy = SmallestChunk(),
+) where {F,Nargs, OutType}
 
     fty = if f <: Enzyme.Annotation
         f.parameters[1]
diff --git a/test/sugar.jl b/test/sugar.jl
index 390e55ed39..0db01594b6 100644
--- a/test/sugar.jl
+++ b/test/sugar.jl
@@ -688,28 +688,40 @@ fchunk2(x) = map(sin, x) + map(cos, reverse(x))
 
     @testset "Forward gradient" begin
         @testset for chunk in strategies
-            for n in (10, 30)
+            for n in (2, 10)
                 x = ones(n)
                 g = gradient(Forward, fchunk1, x)
-                @test g == gradient(Forward, fchunk1, x; chunk)
+                if chunk isa FixedChunk{3} && n == 2
+                    @test_throws ErrorException gradient(Forward, fchunk1, x; chunk)
+                else
+                    @test g == gradient(Forward, fchunk1, x; chunk)
+                end
             end
         end
     end
     @testset "Forward Jacobian" begin
         @testset for chunk in strategies
-            for n in (10, 30)
+            for n in (2, 10)
                 x = ones(n)
                 J = jacobian(Forward, fchunk2, x)
-                @test J == jacobian(Forward, fchunk2, x; chunk)
+                if chunk isa FixedChunk{3} && n == 2
+                    @test_throws ErrorException jacobian(Forward, fchunk2, x; chunk)
+                else
+                    @test J == jacobian(Forward, fchunk2, x; chunk)
+                end
             end
         end
     end
     @testset "Reverse Jacobian" begin
         @testset for chunk in strategies
-            for n in (10, 30)
+            for n in (2, 10)
                 x = ones(n)
                 J = jacobian(Forward, fchunk2, x)
-                @test J == jacobian(Reverse, fchunk2, x; chunk)
+                if chunk isa FixedChunk{3} && n == 2
+                    @test_throws ErrorException jacobian(Reverse, fchunk2, x; chunk)
+                else
+                    @test J == jacobian(Reverse, fchunk2, x; chunk)
+                end
             end
         end
     end

From a08d7a369ea4cb08304754041c23f8e15ae9e459 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Thu, 30 Oct 2025 21:18:15 +0100
Subject: [PATCH 08/10] Fix typos

---
 lib/EnzymeCore/src/EnzymeCore.jl |  2 +-
 test/sugar.jl                    | 46 ++++++++++++++++----------------
 2 files changed, 24 insertions(+), 24 deletions(-)

diff --git a/lib/EnzymeCore/src/EnzymeCore.jl b/lib/EnzymeCore/src/EnzymeCore.jl
index 26e328120c..5994721ae3 100644
--- a/lib/EnzymeCore/src/EnzymeCore.jl
+++ b/lib/EnzymeCore/src/EnzymeCore.jl
@@ -806,7 +806,7 @@ Abstract type gathering strategies for chunk size selection.
 # Subtypes
 
 - [`SmallestChunk`](@ref)
-- [`LargestChunkk`](@ref)
+- [`LargestChunk`](@ref)
 - [`FixedChunk`](@ref)
 - [`AutoChunk`](@ref)
 """
diff --git a/test/sugar.jl b/test/sugar.jl
index 0db01594b6..7d4659848c 100644
--- a/test/sugar.jl
+++ b/test/sugar.jl
@@ -672,29 +672,29 @@ fchunk2(x) = map(sin, x) + map(cos, reverse(x))
 
 @testset "Chunking strategies" begin
     @testset "ChunkedOneHot" begin
-        @test chunkedonehot(ones(3), SmallestChunk()) isa Tuple{NTuple{1},NTuple{1},NTuple{1}}
-        @test chunkedonehot(ones(30), LargestChunk()) isa Tuple{NTuple{30}}
-        @test chunkedonehot(ones(10), LargestChunk()) isa Tuple{NTuple{10}}
-        @test chunkedonehot(ones(30), LargestChunk()) isa Tuple{NTuple{30}}
-        @test chunkedonehot(ones(10), FixedChunk{4}()) isa Tuple{NTuple{4},NTuple{4},NTuple{2}}
-        @test chunkedonehot(ones(10), FixedChunk{5}()) isa Tuple{NTuple{5},NTuple{5}}
-        @test chunkedonehot(ones(10), AutoChunk()) isa Tuple{NTuple{10}}
-        @test chunkedonehot(ones(30), AutoChunk()) isa Tuple{NTuple{16}, NTuple{14}}
-        @test chunkedonehot(ones(30), AutoChunk()) isa Tuple{NTuple{16}, NTuple{14}}
-        @test chunkedonehot(ones(40), AutoChunk()) isa Tuple{NTuple{16}, NTuple{16}, NTuple{8}}
+        @test Enzyme.chunkedonehot(ones(3), Enzyme.SmallestChunk()) isa Tuple{NTuple{1},NTuple{1},NTuple{1}}
+        @test Enzyme.chunkedonehot(ones(30), Enzyme.LargestChunk()) isa Tuple{NTuple{30}}
+        @test Enzyme.chunkedonehot(ones(10), Enzyme.LargestChunk()) isa Tuple{NTuple{10}}
+        @test Enzyme.chunkedonehot(ones(30), Enzyme.LargestChunk()) isa Tuple{NTuple{30}}
+        @test Enzyme.chunkedonehot(ones(10), Enzyme.FixedChunk{4}()) isa Tuple{NTuple{4},NTuple{4},NTuple{2}}
+        @test Enzyme.chunkedonehot(ones(10), Enzyme.FixedChunk{5}()) isa Tuple{NTuple{5},NTuple{5}}
+        @test Enzyme.chunkedonehot(ones(10), Enzyme.AutoChunk()) isa Tuple{NTuple{10}}
+        @test Enzyme.chunkedonehot(ones(30), Enzyme.AutoChunk()) isa Tuple{NTuple{16}, NTuple{14}}
+        @test Enzyme.chunkedonehot(ones(30), Enzyme.AutoChunk()) isa Tuple{NTuple{16}, NTuple{14}}
+        @test Enzyme.chunkedonehot(ones(40), Enzyme.AutoChunk()) isa Tuple{NTuple{16}, NTuple{16}, NTuple{8}}
     end
 
-    strategies = [SmallestChunk(), LargestChunk(), FixedChunk{3}(), AutoChunk()]
+    strategies = [Enzyme.SmallestChunk(), Enzyme.LargestChunk(), Enzyme.FixedChunk{3}(), Enzyme.AutoChunk()]
 
     @testset "Forward gradient" begin
         @testset for chunk in strategies
             for n in (2, 10)
                 x = ones(n)
-                g = gradient(Forward, fchunk1, x)
-                if chunk isa FixedChunk{3} && n == 2
-                    @test_throws ErrorException gradient(Forward, fchunk1, x; chunk)
+                g = Enzyme.gradient(Enzyme.Forward, fchunk1, x)
+                if chunk isa Enzyme.FixedChunk{3} && n == 2
+                    @test_throws ErrorException Enzyme.gradient(Enzyme.Forward, fchunk1, x; chunk)
                 else
-                    @test g == gradient(Forward, fchunk1, x; chunk)
+                    @test g == Enzyme.gradient(Enzyme.Forward, fchunk1, x; chunk)
                 end
             end
         end
@@ -703,11 +703,11 @@ fchunk2(x) = map(sin, x) + map(cos, reverse(x))
         @testset for chunk in strategies
             for n in (2, 10)
                 x = ones(n)
-                J = jacobian(Forward, fchunk2, x)
-                if chunk isa FixedChunk{3} && n == 2
-                    @test_throws ErrorException jacobian(Forward, fchunk2, x; chunk)
+                J = Enzyme.jacobian(Enzyme.Forward, fchunk2, x)
+                if chunk isa Enzyme.FixedChunk{3} && n == 2
+                    @test_throws ErrorException Enzyme.jacobian(Enzyme.Forward, fchunk2, x; chunk)
                 else
-                    @test J == jacobian(Forward, fchunk2, x; chunk)
+                    @test J == Enzyme.jacobian(Enzyme.Forward, fchunk2, x; chunk)
                 end
             end
         end
@@ -716,11 +716,11 @@ fchunk2(x) = map(sin, x) + map(cos, reverse(x))
         @testset for chunk in strategies
             for n in (2, 10)
                 x = ones(n)
-                J = jacobian(Forward, fchunk2, x)
-                if chunk isa FixedChunk{3} && n == 2
-                    @test_throws ErrorException jacobian(Reverse, fchunk2, x; chunk)
+                J = Enzyme.jacobian(Enzyme.Forward, fchunk2, x)
+                if chunk isa Enzyme.FixedChunk{3} && n == 2
+                    @test_throws ErrorException Enzyme.jacobian(Enzyme.Reverse, fchunk2, x; chunk)
                 else
-                    @test J == jacobian(Reverse, fchunk2, x; chunk)
+                    @test J == Enzyme.jacobian(Enzyme.Reverse, fchunk2, x; chunk)
                 end
             end
         end

From bec2f09a2e8c08d31983dbe02332e7923350c85e Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Fri, 21 Nov 2025 08:02:24 +0100
Subject: [PATCH 09/10] Remove SmallestChunk, switch default (nothing) to
 LargestChunk

---
 lib/EnzymeCore/src/EnzymeCore.jl | 38 ++++++--------------------------
 lib/EnzymeCore/test/chunk.jl     |  7 ------
 src/Enzyme.jl                    |  4 ++--
 src/sugar.jl                     | 18 +++++++--------
 test/sugar.jl                    | 22 +++++-------------
 5 files changed, 23 insertions(+), 66 deletions(-)

diff --git a/lib/EnzymeCore/src/EnzymeCore.jl b/lib/EnzymeCore/src/EnzymeCore.jl
index 5994721ae3..9c1bfe1429 100644
--- a/lib/EnzymeCore/src/EnzymeCore.jl
+++ b/lib/EnzymeCore/src/EnzymeCore.jl
@@ -8,7 +8,7 @@ export DefaultABI, FFIABI, InlineABI, NonGenABI
 export BatchDuplicatedFunc
 export within_autodiff, ignore_derivatives
 export needs_primal
-export ChunkStrategy, SmallestChunk, LargestChunk, FixedChunk, AutoChunk, pick_chunksize
+export ChunkStrategy, LargestChunk, FixedChunk, AutoChunk, pick_chunksize
 
 function batch_size end
 
@@ -805,7 +805,6 @@ Abstract type gathering strategies for chunk size selection.
 
 # Subtypes
 
-- [`SmallestChunk`](@ref)
 - [`LargestChunk`](@ref)
 - [`FixedChunk`](@ref)
 - [`AutoChunk`](@ref)
@@ -813,20 +812,13 @@ Abstract type gathering strategies for chunk size selection.
 abstract type ChunkStrategy end
 
 """
-    SmallestChunk()
+    LargestChunk()
 
-Select chunk size equal to 1, so that the corresponding array is processed in as many chunks as it has elements.
+Select chunk size equal to the number of elements, so that the corresponding array is processed in a single chunk.
 
 !!! tip
     In the current Enzyme interface, this strategy is equivalent to setting `chunk = nothing`.
 """
-struct SmallestChunk <: ChunkStrategy end
-
-"""
-    LargestChunk()
-
-Select chunk size equal to the number of elements, so that the corresponding array is processed in a single chunk.
-"""
 struct LargestChunk <: ChunkStrategy end
 
 """
@@ -836,9 +828,6 @@ Select chunk size equal to a fixed integer `C`.
 
 !!! tip
     In the current Enzyme interface, this chunk strategy is equivalent to setting `chunk = Val(C)`.
-
-!!! warning
-    This chunk strategy will error if the corresponding array has length `< C`.
 """
 struct FixedChunk{C} <: ChunkStrategy end
 
@@ -855,10 +844,10 @@ const DEFAULT_CHUNK_SIZE = 16
     pick_chunksize(s::ChunkStrategy, n::Integer)
     pick_chunksize(s::ChunkStrategy, a::AbstractArray)
 
-Compute tLargestChunkze chosen by strategy `s` based on the integer `n` or the array `a` (`n` corresponds to the array's length)
+Compute the chunk size chosen by strategy `s` based on the integer `n` or the array `a` (`n` corresponds to the array's length)
 Return a `Val{C}` object.
-LargestChunk
-- In forward-modeLargestChunkand Jacobians, `a` would be the input array.
+
+- In forward-mode Jacobians, `a` would be the input array.
 - In reverse-mode Jacobians, `a` would be the output array.
 
 !!! warning
@@ -866,24 +855,11 @@ LargestChunk
 """
 function pick_chunksize end
 
-pick_chunksize(::SmallestChunk, a_or_n::Union{Integer,AbstractArray}) = Val(1)
-
 pick_chunksize(::LargestChunk, n::Integer) = Val(n)
 pick_chunksize(::LargestChunk, a::AbstractArray) = Val(length(a))  # allows inference on static arrays
 
 pick_chunksize(::AutoChunk, n::Integer) = Val(min(DEFAULT_CHUNK_SIZE, n))  # TODO: improve
 pick_chunksize(s::AutoChunk, a::AbstractArray) = pick_chunksize(s, length(a))
-
-function pick_chunksize(s::FixedChunk{C}, a_or_n::Union{Integer,AbstractArray}) where {C}
-    check_length(s, a_or_n)
-    return Val{C}()
-end
-
-function check_length(::FixedChunk{C}, n::Integer) where {C}
-    if n < C
-        error("Chunk size $C is larger than length $n")
-    end
-end
-check_length(s::FixedChunk{C}, a::AbstractArray) where {C} = check_length(s, length(a))
+pick_chunksize(::FixedChunk{C}, ::Union{Integer,AbstractArray}) where {C} = Val{C}()
 
 end # module EnzymeCore
diff --git a/lib/EnzymeCore/test/chunk.jl b/lib/EnzymeCore/test/chunk.jl
index a3aed39f3a..45d255e001 100644
--- a/lib/EnzymeCore/test/chunk.jl
+++ b/lib/EnzymeCore/test/chunk.jl
@@ -1,13 +1,6 @@
 using Test
 using EnzymeCore
 
-@testset "SmallestChunk" begin
-    @test pick_chunksize(SmallestChunk(), 10) == Val(1)
-    @test pick_chunksize(SmallestChunk(), ones(10)) == Val(1)
-    @test pick_chunksize(SmallestChunk(), 100) == Val(1)
-    @test pick_chunksize(SmallestChunk(), ones(100)) == Val(1)
-end
-
 @testset "LargestChunk" begin
     @test pick_chunksize(LargestChunk(), 10) == Val(10)
     @test pick_chunksize(LargestChunk(), ones(10)) == Val(10)
diff --git a/src/Enzyme.jl b/src/Enzyme.jl
index 125ac4c57a..b790fae3e6 100644
--- a/src/Enzyme.jl
+++ b/src/Enzyme.jl
@@ -111,11 +111,11 @@ export autodiff,
     make_zero!,
     remake_zero!
 
-import EnzymeCore: ChunkStrategy, SmallestChunk, LargestChunk, FixedChunk, AutoChunk, pick_chunksize
+import EnzymeCore: ChunkStrategy, LargestChunk, FixedChunk, AutoChunk, pick_chunksize
 
 export jacobian, gradient, gradient!, hvp, hvp!, hvp_and_gradient!
 export batch_size, onehot, chunkedonehot
-export SmallestChunk, LargestChunk, FixedChunk, AutoChunk
+export LargestChunk, FixedChunk, AutoChunk
 
 using LinearAlgebra
 import SparseArrays
diff --git a/src/sugar.jl b/src/sugar.jl
index 488ece6717..319f937e8e 100644
--- a/src/sugar.jl
+++ b/src/sugar.jl
@@ -422,10 +422,10 @@ get_strategy(chunk::Type{CS}) where {CS<:ChunkStrategy} = chunk
 
 function get_strategy(::Type{Nothing})
     Base.depwarn(
-        "The `chunk=nothing` configuration will be deprecated in a future release. Please use `chunk=SmallestChunk()` instead.",
+        "The `chunk=nothing` configuration will be deprecated in a future release. Please use `chunk=LargestChunk()` instead.",
         :get_strategy,
     )
-    return SmallestChunk
+    return LargestChunk()
 end
 
 function get_strategy(::Type{Val{C}}) where {C}
@@ -472,7 +472,7 @@ end
             push!(exprs, :(nothing))
         elseif ty <: AbstractFloat
             push!(exprs, :(nothing))
-        elseif chunk_strategy == SmallestChunk || chunk_strategy == FixedChunk{1}
+        elseif chunk_strategy == LargestChunk || chunk_strategy == FixedChunk{1}
             push!(exprs, :(onehot($arg)))
         else
             push!(exprs, :(chunkedonehot($arg, chunk)))
@@ -528,7 +528,7 @@ end
 @inline specialize_output(output, input) = output
 
 """
-    gradient(::ForwardMode, f, x, args...; chunk=SmallestChunk(), shadows=create_shadows(chunk, x, args...))
+    gradient(::ForwardMode, f, x, args...; chunk=LargestChunk(), shadows=create_shadows(chunk, x, args...))
 
 Compute the gradient of an array-input function `f` using forward mode.
 The optional keyword argument `chunk` denotes the chunk size to use: it can be any instance of [`EnzymeCore.ChunkStrategy`](@ref EnzymeCore.ChunkStrategy).
@@ -614,7 +614,7 @@ gradient(Forward, mul, [2.0, 3.0], Const([2.7, 3.1]))
     f::F,
     x::ty_0,
     args::Vararg{Any,N};
-    chunk::ExtendedChunkStrategy = SmallestChunk(),
+    chunk::ExtendedChunkStrategy = LargestChunk(),
     shadows::ST = create_shadows(chunk, x, args...),
 ) where {F, ReturnPrimal,ABI,ErrIfFuncWritten,RuntimeActivity,StrongZero,ST, ty_0, N}
 
@@ -688,7 +688,7 @@ gradient(Forward, mul, [2.0, 3.0], Const([2.7, 3.1]))
             :($resp[1])
         elseif argnum == 0
             vals[i]
-        elseif chunk_strategy == SmallestChunk
+        elseif chunk_strategy == LargestChunk
             dargs = Union{Symbol,Expr}[]
             for (j, arg2) in enumerate(syms)
                 if i == j
@@ -968,7 +968,7 @@ end
         else
             push!(exprs, Expr(:(=), mdi, :(Compiler.active_reg_nothrow($xti) == Compiler.ActiveState || Compiler.active_reg_nothrow($xti) == Compiler.MixedState)))
 
-            if chunk_strategy == SmallestChunk || chunk_strategy == FixedChunk{1}
+            if chunk_strategy == FixedChunk{1}
                 push!(MDTys, :($mdi ? MixedDuplicated{$xti} : Duplicated{$xti}))
             else
                 push!(MDTys, :($mdi ? BatchMixedDuplicated{$xti, $chunksize} : BatchDuplicated{$xti, $chunksize}))
@@ -1197,7 +1197,7 @@ end
 end
 
 """
-    jacobian(::ReverseMode, f, x; n_outs=nothing, chunk=SmallestChunk())
+    jacobian(::ReverseMode, f, x; n_outs=nothing, chunk=LargestChunk())
     jacobian(::ReverseMode, f, x)
 
 Compute the jacobian of a array-output function `f` using (potentially vector) reverse mode.
@@ -1255,7 +1255,7 @@ this function will retun an AbstractArray of shape `size(output)` of values of t
     f::F,
     xs::Vararg{Any, Nargs};
     n_outs::OutType = nothing,
-    chunk::ExtendedChunkStrategy = SmallestChunk(),
+    chunk::ExtendedChunkStrategy = LargestChunk(),
 ) where {F,Nargs, OutType}
 
     fty = if f <: Enzyme.Annotation
diff --git a/test/sugar.jl b/test/sugar.jl
index 7d4659848c..00d13f18fb 100644
--- a/test/sugar.jl
+++ b/test/sugar.jl
@@ -672,10 +672,10 @@ fchunk2(x) = map(sin, x) + map(cos, reverse(x))
 
 @testset "Chunking strategies" begin
     @testset "ChunkedOneHot" begin
-        @test Enzyme.chunkedonehot(ones(3), Enzyme.SmallestChunk()) isa Tuple{NTuple{1},NTuple{1},NTuple{1}}
         @test Enzyme.chunkedonehot(ones(30), Enzyme.LargestChunk()) isa Tuple{NTuple{30}}
         @test Enzyme.chunkedonehot(ones(10), Enzyme.LargestChunk()) isa Tuple{NTuple{10}}
         @test Enzyme.chunkedonehot(ones(30), Enzyme.LargestChunk()) isa Tuple{NTuple{30}}
+        @test Enzyme.chunkedonehot(ones(3), Enzyme.FixedChunk{1}()) isa Tuple{NTuple{1},NTuple{1},NTuple{1}}
         @test Enzyme.chunkedonehot(ones(10), Enzyme.FixedChunk{4}()) isa Tuple{NTuple{4},NTuple{4},NTuple{2}}
         @test Enzyme.chunkedonehot(ones(10), Enzyme.FixedChunk{5}()) isa Tuple{NTuple{5},NTuple{5}}
         @test Enzyme.chunkedonehot(ones(10), Enzyme.AutoChunk()) isa Tuple{NTuple{10}}
@@ -684,18 +684,14 @@ fchunk2(x) = map(sin, x) + map(cos, reverse(x))
         @test Enzyme.chunkedonehot(ones(40), Enzyme.AutoChunk()) isa Tuple{NTuple{16}, NTuple{16}, NTuple{8}}
     end
 
-    strategies = [Enzyme.SmallestChunk(), Enzyme.LargestChunk(), Enzyme.FixedChunk{3}(), Enzyme.AutoChunk()]
+    strategies = [Enzyme.LargestChunk(), Enzyme.FixedChunk{1}(), Enzyme.FixedChunk{3}(), Enzyme.AutoChunk()]
 
     @testset "Forward gradient" begin
         @testset for chunk in strategies
             for n in (2, 10)
                 x = ones(n)
                 g = Enzyme.gradient(Enzyme.Forward, fchunk1, x)
-                if chunk isa Enzyme.FixedChunk{3} && n == 2
-                    @test_throws ErrorException Enzyme.gradient(Enzyme.Forward, fchunk1, x; chunk)
-                else
-                    @test g == Enzyme.gradient(Enzyme.Forward, fchunk1, x; chunk)
-                end
+                @test g == Enzyme.gradient(Enzyme.Forward, fchunk1, x; chunk)
             end
         end
     end
@@ -704,11 +700,7 @@ fchunk2(x) = map(sin, x) + map(cos, reverse(x))
             for n in (2, 10)
                 x = ones(n)
                 J = Enzyme.jacobian(Enzyme.Forward, fchunk2, x)
-                if chunk isa Enzyme.FixedChunk{3} && n == 2
-                    @test_throws ErrorException Enzyme.jacobian(Enzyme.Forward, fchunk2, x; chunk)
-                else
-                    @test J == Enzyme.jacobian(Enzyme.Forward, fchunk2, x; chunk)
-                end
+                @test J == Enzyme.jacobian(Enzyme.Forward, fchunk2, x; chunk)
             end
         end
     end
@@ -717,11 +709,7 @@ fchunk2(x) = map(sin, x) + map(cos, reverse(x))
             for n in (2, 10)
                 x = ones(n)
                 J = Enzyme.jacobian(Enzyme.Forward, fchunk2, x)
-                if chunk isa Enzyme.FixedChunk{3} && n == 2
-                    @test_throws ErrorException Enzyme.jacobian(Enzyme.Reverse, fchunk2, x; chunk)
-                else
-                    @test J == Enzyme.jacobian(Enzyme.Reverse, fchunk2, x; chunk)
-                end
+                @test J == Enzyme.jacobian(Enzyme.Reverse, fchunk2, x; chunk)
             end
         end
     end

From 9e18f0366eab82addaeeff6d47d0473b617c9df2 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Fri, 21 Nov 2025 17:37:19 +0100
Subject: [PATCH 10/10] Add constructor

---
 lib/EnzymeCore/src/EnzymeCore.jl | 3 +++
 lib/EnzymeCore/test/chunk.jl     | 1 +
 2 files changed, 4 insertions(+)

diff --git a/lib/EnzymeCore/src/EnzymeCore.jl b/lib/EnzymeCore/src/EnzymeCore.jl
index 9c1bfe1429..37d96040f4 100644
--- a/lib/EnzymeCore/src/EnzymeCore.jl
+++ b/lib/EnzymeCore/src/EnzymeCore.jl
@@ -823,6 +823,7 @@ struct LargestChunk <: ChunkStrategy end
 
 """
     FixedChunk{C}()
+    FixedChunk(C)  # type-unstable
 
 Select chunk size equal to a fixed integer `C`.
 
@@ -831,6 +832,8 @@ Select chunk size equal to a fixed integer `C`.
 """
 struct FixedChunk{C} <: ChunkStrategy end
 
+FixedChunk(C::Int) = FixedChunk{C}()
+
 """
     AutoChunk()
 
diff --git a/lib/EnzymeCore/test/chunk.jl b/lib/EnzymeCore/test/chunk.jl
index 45d255e001..f4e56612a6 100644
--- a/lib/EnzymeCore/test/chunk.jl
+++ b/lib/EnzymeCore/test/chunk.jl
@@ -9,6 +9,7 @@ using EnzymeCore
 end
 
 @testset "FixedChunk" begin
+    @test FixedChunk(3) == FixedChunk{3}()
     @test_throws ErrorException pick_chunksize(FixedChunk{3}(), 2)
     @test_throws ErrorException pick_chunksize(FixedChunk{3}(), ones(2))
     @test pick_chunksize(FixedChunk{3}(), 10) == Val(3)