From 99bc756e0ee0277654677da081c78d4c1ce0b6dd Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Tue, 2 Jul 2024 22:36:56 -0700
Subject: [PATCH] test: skip certain RNG tests for cuda/rocm

---
 Project.toml               | 3 ++-
 test/initializers_tests.jl | 7 ++++++-
 test/runtests.jl           | 1 -
 test/shared_testsetup.jl   | 6 +++---
 4 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/Project.toml b/Project.toml
index ca2b7f0..e66ab80 100644
--- a/Project.toml
+++ b/Project.toml
@@ -52,10 +52,11 @@ oneAPI = "1.5.0"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
+GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "Documenter", "ExplicitImports", "Pkg", "ReTestItems", "StableRNGs", "Test"]
+test = ["Aqua", "Documenter", "ExplicitImports", "GPUArrays", "Pkg", "ReTestItems", "StableRNGs", "Test"]
diff --git a/test/initializers_tests.jl b/test/initializers_tests.jl
index c6e1818..af968f8 100644
--- a/test/initializers_tests.jl
+++ b/test/initializers_tests.jl
@@ -272,7 +272,12 @@ end
 
                 v = kaiming_normal(rng, n_in, n_out)
                 σ2 = sqrt(2 / n_out)
-                @test 0.9σ2 < std(Array(v)) < 1.1σ2 # Just for safety move to Array
+
+                if (backend == "cuda" || backend == "amdgpu") && rng isa GPUArrays.RNG
+                    @test_broken 0.9σ2 < std(v) < 1.1σ2
+                else
+                    @test 0.9σ2 < std(v) < 1.1σ2
+                end
             end
             # Type
             @test eltype(kaiming_uniform(rng, 3, 4; gain=1.5f0)) == Float32
diff --git a/test/runtests.jl b/test/runtests.jl
index db4d5e8..994df2b 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -8,7 +8,6 @@ BACKEND_GROUP == "all" || BACKEND_GROUP == "cuda" && push!(EXTRA_PKGS, "CUDA")
 BACKEND_GROUP == "all" || BACKEND_GROUP == "amdgpu" && push!(EXTRA_PKGS, "AMDGPU")
 BACKEND_GROUP == "all" || BACKEND_GROUP == "metal" && push!(EXTRA_PKGS, "Metal")
 BACKEND_GROUP == "all" || BACKEND_GROUP == "oneapi" && push!(EXTRA_PKGS, "oneAPI")
-BACKEND_GROUP != "all" && push!(EXTRA_PKGS, "GPUArrays")
 
 if !isempty(EXTRA_PKGS)
     @info "Installing Extra Packages for testing" EXTRA_PKGS=EXTRA_PKGS
diff --git a/test/shared_testsetup.jl b/test/shared_testsetup.jl
index e3461ba..8d7cb83 100644
--- a/test/shared_testsetup.jl
+++ b/test/shared_testsetup.jl
@@ -1,6 +1,6 @@
 @testsetup module SharedTestSetup
 
-using GPUArraysCore, Random, StableRNGs
+using GPUArrays, GPUArraysCore, Random, StableRNGs
 
 GPUArraysCore.allowscalar(false)
 
@@ -13,7 +13,7 @@ if BACKEND_GROUP == "all" || BACKEND_GROUP == "cpu"
             (Random.GLOBAL_RNG, AbstractArray, true, "cpu")])
 end
 if BACKEND_GROUP == "all" || BACKEND_GROUP == "cuda"
-    using CUDA, GPUArrays
+    using CUDA
     append!(RNGS_ARRTYPES,
         [(CUDA.default_rng(), CuArray, true, "cuda"),
             (GPUArrays.default_rng(CuArray), CuArray, true, "cuda")])
@@ -38,6 +38,6 @@ if BACKEND_GROUP == "all" || BACKEND_GROUP == "oneapi"
     push!(RNGS_ARRTYPES, (oneAPI.gpuarrays_rng(), oneArray, supports_fp64, "oneapi"))
 end
 
-export StableRNG, RNGS_ARRTYPES, BACKEND_GROUP
+export StableRNG, RNGS_ARRTYPES, BACKEND_GROUP, GPUArrays
 
 end