feat: remove closure restriction on batching (EnzymeAD#1443)

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "Reactant"
 uuid = "3c362404-f566-11ee-1572-e11a4b42c853"
 authors = ["William Moses <[email protected]>", "Valentin Churavy <[email protected]>", "Sergio Sánchez Ramírez <[email protected]>", "Paul Berg <[email protected]>", "Avik Pal <[email protected]>", "Mosè Giordano <[email protected]>"]
-version = "0.2.142"
+version = "0.2.143"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/Ops.jl

@@ -2876,7 +2876,6 @@ end
         args_in_result=:none,
         do_transpose=false,
     )
-    @assert !mlir_fn_res.fnwrapped "Currently we don't support batching closures."
 
     func = mlir_fn_res.f
     @assert MLIR.IR.nregions(func) == 1

test/batching.jl

@@ -2,6 +2,12 @@ using Reactant, Test
 
 f1(x::AbstractMatrix) = sum(x; dims=1)
 
+f2(x::AbstractMatrix, y::Int) = x .+ y
+
+function f3(x::AbstractArray{T,3}, y::Int) where {T}
+    return Reactant.Ops.batch(Base.Fix2(f2, y), x, [1, 2])
+end
+
 @testset "mapslices" begin
     A = collect(reshape(1:30, (2, 5, 3)))
     A_ra = Reactant.to_rarray(A)
@@ -16,3 +22,10 @@ f1(x::AbstractMatrix) = sum(x; dims=1)
 
     @test B ≈ B_ra
 end
+
+@testset "closure" begin
+    A = collect(reshape(1:30, (2, 5, 3)))
+    A_ra = Reactant.to_rarray(A)
+
+    @test @jit(f3(A_ra, 1)) ≈ A .+ 1
+end