wrap comm test

Author: jumerckx

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