test: fixes

Author: avik-pal

ext/ReactantFillArraysExt.jl

@@ -1,6 +1,6 @@
 module ReactantFillArraysExt
 
-using Reactant: Reactant, TracedUtils, TracedRNumber, Ops, Sharding, unwrapped_eltype
+using Reactant: Reactant, TracedRNumber, Sharding, unwrapped_eltype
 using ReactantCore: ReactantCore
 using FillArrays: FillArrays, AbstractFill, Fill, Ones, Zeros, OneElement
 using GPUArraysCore: @allowscalar

test/core/wrapped_arrays.jl

@@ -1,4 +1,4 @@
-using Reactant, Test, Statistics, NNlib, LinearAlgebra
+using Reactant, Test, Statistics, LinearAlgebra
 
 function view_getindex_1(x)
     x = view(x, 2:3, 1:2, :)
@@ -77,18 +77,6 @@ end
     @test v1 ≈ v2
 end
 
-function btranspose_badjoint(x)
-    x1 = NNlib.batched_transpose(x)
-    x2 = NNlib.batched_adjoint(x)
-    return x1 .+ x2
-end
-
-@testset "batched transpose/adjoint" begin
-    x = rand(4, 2, 3)
-    x_ra = Reactant.to_rarray(x)
-    @test @jit(btranspose_badjoint(x_ra)) ≈ btranspose_badjoint(x)
-end
-
 function bypass_permutedims(x)
     x = PermutedDimsArray(x, (2, 1, 3))  # Don't use permutedims here
     return view(x, 2:3, 1:2, :)

test/integration/NNlib/runtests.jl

@@ -792,3 +792,15 @@ end
     hlo = repr(@code_hlo(NNlib.gather(x_ra, idxs_ra)))
     @test !contains(hlo, "i64>")
 end
+
+function btranspose_badjoint(x)
+    x1 = NNlib.batched_transpose(x)
+    x2 = NNlib.batched_adjoint(x)
+    return x1 .+ x2
+end
+
+@testset "batched transpose/adjoint" begin
+    x = rand(4, 2, 3)
+    x_ra = Reactant.to_rarray(x)
+    @test @jit(btranspose_badjoint(x_ra)) ≈ btranspose_badjoint(x)
+end