fix test GPU

Author: glou-nes

src/Ops.jl

@@ -961,13 +961,13 @@ end
     @assert 0 < dimension <= ndims(x) "$x invalid dimension"
 
     (a, b) = (Reactant.ConcreteRNumber(T(0)), Reactant.ConcreteRNumber(T(0)))
-    func = Reactant.TracedUtils.make_mlir_fn(comparator, (a, b), (), "main"; no_args_in_result=true, return_dialect=:stablehlo)[2]
+    func = Reactant.TracedUtils.make_mlir_fn(comparator, (a, b), (), "comparator"; no_args_in_result=true, return_dialect=:stablehlo)[2]
     @assert MLIR.IR.nregions(func) == 1
     fn_name = String(
         MLIR.IR.attr(func, String(MLIR.API.mlirSymbolTableGetSymbolAttributeName()))
     )
     #C5:
-    @assert fn_name == "main" "$comparator: no function generated"
+    @assert fn_name == "comparator" "$comparator: no function generated"
     ftype_attr = MLIR.IR.attr(func, "function_type")
     ftype = MLIR.IR.Type(ftype_attr)
     @assert MLIR.IR.result(ftype) == MLIR.IR.TensorType((), MLIR.IR.Type(Bool)) error(

test/ops.jl

@@ -692,10 +692,11 @@ end
     basic_sort(x, dimension) = Reactant.Ops.sort(x; comparator=(a, b) -> a < b, dimension)
     for i in 1:3
         t_size = tuple(fill(10, (i,))...)
-        x = Reactant.to_rarray(randn(t_size))
+        x = randn(t_size)
+        xa = Reactant.to_rarray(x)
 
         for j in 1:i
-            @test (i == 1 ? sort(x) : sort(x; dims=j)) == @jit basic_sort(x, j)
+            @test (i == 1 ? sort(x) : sort(x; dims=j)) == @jit basic_sort(xa, j)
         end
     end
 end