Add rand! implementation

Author: pszufe

Project.toml

@@ -28,11 +28,13 @@ UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 [weakdeps]
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 GraphViz = "f526b714-d49f-11e8-06ff-31ed36ee7ee0"
 JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 
 [extensions]
+DistributionsExt = "Distributions"
 GraphVizExt = "GraphViz"
 GraphVizSimpleExt = "Colors"
 JSON3Ext = "JSON3"
@@ -43,6 +45,7 @@ Adapt = "4.0.4"
 Colors = "0.12"
 DataFrames = "1"
 DataStructures = "0.18"
+Distributions = "0.25"
 GraphViz = "0.2"
 Graphs = "1"
 MacroTools = "0.5"
@@ -61,6 +64,7 @@ julia = "1.8"
 [extras]
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 GraphViz = "f526b714-d49f-11e8-06ff-31ed36ee7ee0"
 JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"

docs/src/darray.md

@@ -431,6 +431,10 @@ From `Base`:
 - `map`/`reduce`/`mapreduce`
 - `sum`/`prod`
 - `minimum`/`maximum`/`extrema`
+- `map!`
+
+From `Random`:
+- `rand!`/`randn!`
 
 From `Statistics`:
 - `mean`

ext/DistributionsExt.jl

@@ -0,0 +1,15 @@
+module DistributionsExt
+
+if isdefined(Base, :get_extension)
+    using Distributions
+else
+    using ..Distributions
+end
+
+using Dagger, Random
+
+Random.rand!(s::Sampleable, A::DArray{T}) where T = map!(_ -> rand(s), A)
+Random.rand!(rng::AbstractRNG, s::Sampleable{Univariate}, A::DArray{T}) where T = map!(_ -> rand(rng, s), A)
+Random.rand!(rng::AbstractRNG, s::Sampleable{ArrayLikeVariate{M}}, A::DArray{T}) where {M, T} = map!(_ -> rand(rng, s), A)
+
+end # module DistributionsExt

src/Dagger.jl

@@ -13,6 +13,8 @@ import Distributed: Future, RemoteChannel, myid, workers, nworkers, procs, remot
 
 import LinearAlgebra
 import LinearAlgebra: Adjoint, BLAS, Diagonal, Bidiagonal, Tridiagonal, LAPACK, LowerTriangular, PosDefException, Transpose, UpperTriangular, UnitLowerTriangular, UnitUpperTriangular, diagind, ishermitian, issymmetric
+import Random
+import Random: AbstractRNG
 
 import UUIDs: UUID, uuid4
 
@@ -82,6 +84,7 @@ include("array/sort.jl")
 include("array/linalg.jl")
 include("array/mul.jl")
 include("array/cholesky.jl")
+include("array/random.jl")
 
 # Visualization
 include("visualization.jl")
@@ -101,6 +104,9 @@ function __init__()
     system_uuid()
 
     @static if !isdefined(Base, :get_extension)
+        @require Distributions="31c24e10-a181-5473-b8eb-7969acd0382f" begin
+            include(joinpath(dirname(@__DIR__), "ext", "DistributionsExt.jl"))
+        end
         @require Graphs="86223c79-3864-5bf0-83f7-82e725a168b6" begin
             @require GraphViz="f526b714-d49f-11e8-06ff-31ed36ee7ee0" begin
                 include(joinpath(dirname(@__DIR__), "ext", "GraphVizExt.jl"))

src/array/operators.jl

@@ -111,3 +111,24 @@ end
 
 Base.first(A::DArray) = A[begin]
 Base.last(A::DArray) = A[end]
+
+# In-place operations
+
+function Base.map!(f, a::DArray{T}) where T
+    Dagger.spawn_datadeps() do
+        for ca in chunks(a)
+            Dagger.@spawn map!(f, InOut(ca), ca)
+        end
+    end
+    return a
+end
+
+function Base.map!(f, a::DArray{T}, b::AbstractArray{U}) where {T, U}
+    b2 = view(b, a.partitioning)
+    Dagger.spawn_datadeps() do
+        for (c_a, c_b2) in zip(chunks(a), chunks(b2))
+            Dagger.@spawn map!(f, InOut(c_a), c_b2)
+        end
+    end
+    return a
+end

src/array/random.jl

@@ -0,0 +1,26 @@
+Random.rand!(dm::DArray{T}) where T = map!(_ -> rand(T), dm)
+Random.randn!(dm::DArray{T}) where T = map!(_ -> randn(T), dm)
+
+randfork(rng::R, n::Integer) where {R<:AbstractRNG} =
+    R(abs(rand(copy(rng), Int) + n))
+
+function Random.rand!(rng::AbstractRNG, A::DArray{T}) where T
+    part_sz = prod(map(length, first(A.subdomains).indexes))
+    Dagger.spawn_datadeps() do
+        for Ac in chunks(A)
+            rng = randfork(rng, part_sz)
+            Dagger.@spawn map!(_->rand(rng, T), InOut(Ac), Ac)
+        end
+    end
+    return A
+end
+function Random.randn!(rng::AbstractRNG, A::DArray{T}) where T
+    part_sz = prod(map(length, first(A.subdomains).indexes))
+    Dagger.spawn_datadeps() do
+        for Ac in chunks(A)
+            rng = randfork(rng, part_sz)
+            Dagger.@spawn map!(_->randn(rng, T), InOut(Ac), Ac)
+        end
+    end
+    return A
+end

test/Project.toml

@@ -3,8 +3,9 @@ ArgParse = "c7e460c6-2fb9-53a9-8c5b-16f535851c63"
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
-Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 GraphViz = "f526b714-d49f-11e8-06ff-31ed36ee7ee0"
+Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MemPool = "f9f48841-c794-520a-933b-121f7ba6ed94"

test/array/core.jl

@@ -12,6 +12,16 @@ end
     @test collect(X1) .+ 1 == collect(X2)
 end
 
+@testset "map!" begin
+    X = zeros(Blocks(10,10), Int, 40, 40)
+
+    map!(_ -> 1, X)
+    @test sum(X) == length(X)
+
+    map!(x -> 2x+1, X, Matrix(X))
+    @test sum(X) == 3*length(X)
+end
+
 @testset "DiffEq support" begin
     X = ones(Blocks(10), 100)
     X0 = zero(X)

test/array/random.jl

@@ -0,0 +1,30 @@
+using Distributions, Statistics
+
+@testset "rand!" begin
+    X = zeros(Blocks(10,10), Float64, 40, 40)
+
+    Random.rand!(X)
+    @test all(0.0 .<= X .<= 1.0)
+    @test sum(X) > 0.01
+
+    Random.rand!(Uniform(1, 10), X)
+    @test all(1.0 .<= X .<= 10.0)
+    @test sum(X) > length(X) + 0.1
+
+    for RNGT in (MersenneTwister, Xoshiro)
+        Random.rand!(RNGT(1234), X)
+        @test all(0.0 .<= X .<= 1.0)
+        @test sum(X) > 0.01
+
+        Random.rand!(RNGT(1234), Uniform(1,10), X)
+        @test all(1.0 .<= X .<= 10.0)
+        @test sum(X) > length(X) + 0.1
+
+        Random.randn!(X)
+        @test sum(X) <= 2*length(X)
+
+        Random.randn!(RNGT(1234), X)
+        @test sum(X) <= 2*length(X)
+        @test all(X .!= 0.0)
+    end
+end

test/runtests.jl

@@ -18,6 +18,7 @@ tests = [
     ("Array - MapReduce", "array/mapreduce.jl"),
     ("Array - LinearAlgebra - Matmul", "array/linalg/matmul.jl"),
     ("Array - LinearAlgebra - Cholesky", "array/linalg/cholesky.jl"),
+    ("Array - Random", "array/random.jl"),
     ("Caching", "cache.jl"),
     ("Disk Caching", "diskcaching.jl"),
     ("File IO", "file-io.jl"),