Merge pull request #41 from ChrisRackauckas-Claude/fix-enzyme-batch-width

Generalize Enzyme forward rules to arbitrary batch width

Author: ChrisRackauckas

ext/FunctionWrappersWrappersEnzymeExt.jl

@@ -6,43 +6,54 @@ using EnzymeCore
 using EnzymeCore.EnzymeRules
 
 # =============================================================================
-# Forward mode rules
+# Forward mode rules — generalized to arbitrary batch width W
 # =============================================================================
 
 # Shadow only (Forward mode, no primal)
 function EnzymeRules.forward(
-    ::EnzymeRules.FwdConfig{false, true, 1, RuntimeActivity, StrongZero},
+    ::EnzymeRules.FwdConfig{false, true, W, RuntimeActivity, StrongZero},
     func::EnzymeCore.Const{<:FunctionWrappersWrapper},
     RT::Type{<:EnzymeCore.Annotation{T}},
     args::Vararg{EnzymeCore.Annotation, N}
-) where {T, N, RuntimeActivity, StrongZero}
+) where {T, W, N, RuntimeActivity, StrongZero}
     f_orig = unwrap(func.val)
-    shadow_result = Enzyme.autodiff(Forward, Const(f_orig), Duplicated{T}, args...)
-    return shadow_result[1]::T
+    if W == 1
+        shadow_result = Enzyme.autodiff(Forward, Const(f_orig), Duplicated{T}, args...)
+        return shadow_result[1]::T
+    else
+        shadow_result = Enzyme.autodiff(Forward, Const(f_orig), BatchDuplicated{T, W}, args...)
+        return shadow_result[1]::NTuple{W, T}
+    end
 end
 
 # Both primal and shadow (ForwardWithPrimal mode)
 function EnzymeRules.forward(
-    ::EnzymeRules.FwdConfig{true, true, 1, RuntimeActivity, StrongZero},
+    ::EnzymeRules.FwdConfig{true, true, W, RuntimeActivity, StrongZero},
     func::EnzymeCore.Const{<:FunctionWrappersWrapper},
     RT::Type{<:EnzymeCore.Annotation{T}},
     args::Vararg{EnzymeCore.Annotation, N}
-) where {T, N, RuntimeActivity, StrongZero}
+) where {T, W, N, RuntimeActivity, StrongZero}
     f_orig = unwrap(func.val)
     pargs = ntuple(i -> args[i].val, Val(N))
     primal = f_orig(pargs...)::T
-    shadow_result = Enzyme.autodiff(Forward, Const(f_orig), Duplicated{T}, args...)
-    shadow = shadow_result[1]::T
-    return Duplicated(primal, shadow)
+    if W == 1
+        shadow_result = Enzyme.autodiff(Forward, Const(f_orig), Duplicated{T}, args...)
+        shadow = shadow_result[1]::T
+        return Duplicated(primal, shadow)
+    else
+        shadow_result = Enzyme.autodiff(Forward, Const(f_orig), BatchDuplicated{T, W}, args...)
+        shadows = shadow_result[1]::NTuple{W, T}
+        return BatchDuplicated(primal, shadows)
+    end
 end
 
-# Primal only (Const return type)
+# Primal only (Const return type) — width-independent
 function EnzymeRules.forward(
-    ::EnzymeRules.FwdConfig{true, false, 1, RuntimeActivity, StrongZero},
+    ::EnzymeRules.FwdConfig{true, false, W, RuntimeActivity, StrongZero},
     func::EnzymeCore.Const{<:FunctionWrappersWrapper},
     RT::Type{<:EnzymeCore.Annotation},
     args::Vararg{EnzymeCore.Annotation, N}
-) where {N, RuntimeActivity, StrongZero}
+) where {W, N, RuntimeActivity, StrongZero}
     f_orig = unwrap(func.val)
     pargs = ntuple(i -> args[i].val, Val(N))
     return f_orig(pargs...)

test/enzyme_tests.jl

@@ -50,3 +50,27 @@ end
     result = Enzyme.autodiff(Reverse, Const(fww_sin), Active, Active(1.0))
     @test result[1][1] ≈ cos(1.0)
 end
+
+@testset "Enzyme batch forward mode (width > 1)" begin
+    f(x) = x^2
+    fww = FunctionWrappersWrapper(f, (Tuple{Float64},), (Float64,))
+
+    # Batch width = 2: compute derivatives for two tangent directions simultaneously.
+    # f(x) = x^2 → f'(x) = 2x; at x=3.0 with tangents (1.0, 2.0) → shadows (6.0, 12.0)
+    result = Enzyme.autodiff(
+        Forward, Const(fww), BatchDuplicated,
+        BatchDuplicated(3.0, (1.0, 2.0))
+    )
+    shadows = result[1]
+    @test shadows[1] ≈ 6.0   # f'(3) * 1.0
+    @test shadows[2] ≈ 12.0  # f'(3) * 2.0
+
+    # ForwardWithPrimal, batch width = 2
+    result_wp = Enzyme.autodiff(
+        ForwardWithPrimal, Const(fww), BatchDuplicated,
+        BatchDuplicated(3.0, (1.0, 2.0))
+    )
+    @test result_wp[1][1] ≈ 6.0   # shadow 1
+    @test result_wp[1][2] ≈ 12.0  # shadow 2
+    @test result_wp[2] ≈ 9.0      # primal f(3) = 9
+end