fix: pass in RNG to shuffle (#1188)

LuxDL · Jan 8, 2025 · 76dfdb5 · 76dfdb5 · avik-pal · Jan 8, 2025
1 parent d2dc63f
commit 76dfdb5
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 3 deletions.
diff --git a/lib/WeightInitializers/Project.toml b/lib/WeightInitializers/Project.toml
@@ -1,7 +1,7 @@
 name = "WeightInitializers"
 uuid = "d49dbf32-c5c2-4618-8acc-27bb2598ef2d"
 authors = ["Avik Pal <[email protected]> and contributors"]
-version = "1.0.4"
+version = "1.0.5"
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"

diff --git a/lib/WeightInitializers/src/WeightInitializers.jl b/lib/WeightInitializers/src/WeightInitializers.jl
@@ -3,7 +3,7 @@ module WeightInitializers
 using ArgCheck: @argcheck
 using GPUArraysCore: @allowscalar
 using LinearAlgebra: LinearAlgebra, Diagonal, qr
-using Random: Random, AbstractRNG, shuffle
+using Random: Random, AbstractRNG
 using SpecialFunctions: SpecialFunctions, erfinv  # TODO: Move to Ext in v2.0
 using Statistics: Statistics, std
 

diff --git a/lib/WeightInitializers/src/initializers.jl b/lib/WeightInitializers/src/initializers.jl
@@ -225,7 +225,19 @@ function sparse_init(rng::AbstractRNG, ::Type{T}, dims::Integer...;
     sparse_array .*= T(std)
     fill!(view(sparse_array, 1:num_zeros, :), zero(T))
 
-    return @allowscalar mapslices(shuffle, sparse_array; dims=1)
+    if applicable(Random.rng_native_52, rng)
+        @inbounds for i in axes(sparse_array, 2)
+            @allowscalar Random.shuffle!(rng, view(sparse_array, :, i))
+        end
+    else
+        @warn "`rng` is not supported by `Random.shuffle!`. Ignoring the `rng` for \
+               shuffle." maxlog=1
+        @inbounds for i in axes(sparse_array, 2)
+            @allowscalar Random.shuffle!(view(sparse_array, :, i))
+        end
+    end
+
+    return sparse_array
 end
 
 """