Slightly faster rand! on CPU arrays with MPS rng

JuliaGPU · Apr 10, 2024 · 84911a9 · 84911a9
1 parent 2bd3eb3
commit 84911a9
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/lib/mps/random.jl b/lib/mps/random.jl
@@ -75,14 +75,14 @@ function Random.rand!(rng::RNG, A::AbstractArray{T,N}) where {T <: Union{Uniform
     isempty(A) && return A
     B = MtlArray{T,N,Shared}(undef, size(A))
     rand!(rng, B)
-    copyto!(A, B)
+    copyto!(A, unsafe_wrap(Array{T},B))
     return A
 end
 function Random.randn!(rng::RNG, A::AbstractArray{T,N}) where {T <: Float32, N}
     isempty(A) && return A
     B = MtlArray{T,N,Shared}(undef, size(A))
     randn!(rng, B)
-    copyto!(A, B)
+    copyto!(A, unsafe_wrap(Array{T},B))
     return A
 end