Test communication ops for wrap lowering

src/Ops.jl

@@ -3564,12 +3564,16 @@ end
                                               $(size(input, dimension)) (got $(lhs))"
     @assert 0 ≤ rhs ≤ size(input, dimension) "rhs must be between 0 and \
                                               $(size(input, dimension)) (got $(rhs))"
+
+    sz = collect(Int64, size(input))
+    sz[dimension] = sz[dimension] + lhs + rhs
+
     return TracedRArray{T,N}(
         (),
         MLIR.IR.result(
             enzymexla.wrap(input.mlir_data; lhs, rhs, dimension=dimension - 1, location), 1
         ),
-        size(input),
+        sz,
     )
 end
 

test/ops.jl

@@ -1138,6 +1138,15 @@ end
     @test fr!(vr) ≈ f!(v)
 end
 
+fn_test_wrap(x) = Reactant.Ops.wrap(x, 2, 1; dimension=3)
+
+@testset "Ops.wrap" begin
+    x = Reactant.to_rarray(rand(2, 3, 4, 5))
+    out = @jit fn_test_wrap(x)
+
+    @test size(out) == (2, 3, 7, 5)
+end
+
 @testset "Ops.fill" begin
     @testset "Fill with TracedScalar" begin
         fn(x) = Ops.fill(x, [2, 3])

test/optimize_comm.jl

@@ -22,6 +22,10 @@ function dus2(x, y)
     return nothing
 end
 
+function wrap(x)
+    return Reactant.Ops.@opcall wrap(x, 7, 7; dimension=1)
+end
+
 if length(addressable_devices) ≥ 8
     @testset "Rotate" begin
         N = min((length(Reactant.devices()) ÷ 2) * 2, 8)
@@ -108,4 +112,16 @@ if length(addressable_devices) ≥ 8
         @test all(x .== convert(Array, rx))
         @test all(y .== convert(Array, ry))
     end
+
+    @testset "Wrap" begin
+        mesh = Sharding.Mesh(Reactant.devices(), (:x,))
+        sharding = Sharding.NamedSharding(mesh, (:x,))
+
+        x = Reactant.to_rarray(rand(8192); sharding)
+        hlo = repr(@code_xla wrap(x))
+
+        @test !contains(hlo, "all-to-all")
+        @test !contains(hlo, "all-gather")
+        @test contains(hlo, "collective-permute")
+    end
 end