fix: pass in RNG to shuffle

LuxDL · Jan 7, 2025 · b83d8b8 · b83d8b8
1 parent d2dc63f
commit b83d8b8
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 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 <avikpal@mit.edu> and contributors"]
-version = "1.0.4"
+version = "1.0.5"
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"

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
 
 """