fix: indexing for gpu_rand

Author: avik-pal

src/host/random.jl

@@ -97,8 +97,9 @@ function Random.randn!(rng::RNG, A::AnyGPUArray{T}) where T <: Number
     @kernel function randn!(a, randstates)
         i = @index(Global, Linear)
         idx = 2*(i - 1) + 1
-        U1 = gpu_rand(T, i, randstates)
-        U2 = gpu_rand(T, i, randstates)
+        threadidx = (i-1)%length(randstates) + 1
+        U1 = gpu_rand(T, threadidx, randstates)
+        U2 = gpu_rand(T, threadidx, randstates)
         Z0 = sqrt(T(-2.0)*log(U1))*cos(T(2pi)*U2)
         Z1 = sqrt(T(-2.0)*log(U1))*sin(T(2pi)*U2)
         @inbounds a[idx] = Z0

test/testsuite/random.jl

@@ -7,7 +7,7 @@
     cpu_rng = Random.default_rng()
 
     @testset "rand" begin  # uniform
-        for T in eltypes, d in (10, (10,10))
+        for T in eltypes, d in (10, (10,10), (128,128))
             A = AT{T}(undef, d)
             B = copy(A)
             rand!(rng, A)