@@ -368,7 +368,9 @@ Can return a view into `localpart(A)`
368368end
369369
370370# shortcut to set/get localparts of a distributed object
371- function Base. getindex (d:: DArray , s:: Symbol )
371+ Base. getindex (d:: DArray , s:: Symbol ) = _getindex (d, s)
372+ Base. getindex (d:: DArray{<:Any, 1} , s:: Symbol ) = _getindex (d, s)
373+ function _getindex (d:: DArray , s:: Symbol )
372374 @assert s in [:L , :l , :LP , :lp ]
373375 return localpart (d)
374376end
@@ -497,10 +499,28 @@ dfill(v, d1::Integer, drest::Integer...) = dfill(v, convert(Dims, tuple(d1, dres
497499Construct a distributed uniform random array.
498500Trailing arguments are the same as those accepted by `DArray`.
499501"""
500- drand (r, dims:: Dims , args... ) = DArray (I -> rand (r, map (length,I)), dims, args... )
501- drand (r, d1:: Integer , drest:: Integer... ) = drand (r, convert (Dims, tuple (d1, drest... )))
502- drand (d1:: Integer , drest:: Integer... ) = drand (Float64, convert (Dims, tuple (d1, drest... )))
503- drand (d:: Dims , args... ) = drand (Float64, d, args... )
502+ drand (:: Type{T} , dims:: Dims ) where {T} = DArray (I -> rand (T, map (length, I)), dims)
503+ drand (X, dims:: Dims ) = DArray (I -> rand (X, map (length, I)), dims)
504+ drand (dims:: Dims ) = drand (Float64, dims)
505+
506+ drand (:: Type{T} , d1:: Integer , drest:: Integer... ) where {T} = drand (T, Dims ((d1, drest... )))
507+ drand (X, d1:: Integer , drest:: Integer... ) = drand (X, Dims ((d1, drest... )))
508+ drand (d1:: Integer , drest:: Integer... ) = drand (Float64, Dims ((d1, drest... )))
509+
510+ # With optional process IDs and number of chunks
511+ for N in (1 , 2 )
512+ @eval begin
513+ drand (:: Type{T} , dims:: Dims , args:: Vararg{Any,$N} ) where {T} = DArray (I -> rand (T, map (length, I)), dims, args... )
514+ drand (X, dims:: Dims , args:: Vararg{Any,$N} ) = DArray (I -> rand (X, map (length, I)), dims, args... )
515+ drand (dims:: Dims , args:: Vararg{Any,$N} ) = drand (Float64, dims, args... )
516+ end
517+ end
518+
519+ # Fix method ambiguities
520+ drand (dims:: Dims , procs:: Tuple{Vararg{Int}} ) = drand (Float64, dims, procs)
521+ drand (dims:: Dims , procs:: Tuple{Vararg{Int}} , dist) = drand (Float64, dims, procs, dist)
522+ drand (X:: Tuple{Vararg{Int}} , dim:: Integer ) = drand (X, Dims ((dim,)))
523+ drand (X:: Tuple{Vararg{Int}} , d1:: Integer , d2:: Integer ) = drand (X, Dims ((d1, d2)))
504524
505525"""
506526 drandn(dims, ...)
@@ -620,7 +640,7 @@ allowscalar(flag = true) = (_allowscalar[] = flag)
620640_scalarindexingallowed () = _allowscalar[] || throw (ErrorException (" scalar indexing disabled" ))
621641
622642getlocalindex (d:: DArray , idx... ) = localpart (d)[idx... ]
623- function getindex_tuple (d:: DArray{T} , I:: Tuple{Vararg{ Int}} ) where T
643+ function getindex_tuple (d:: DArray{T,N } , I:: NTuple{N, Int} ) where {T,N}
624644 chidx = locate (d, I... )
625645 idxs = d. indices[chidx... ]
626646 localidx = ntuple (i -> (I[i] - first (idxs[i]) + 1 ), ndims (d))
@@ -632,11 +652,10 @@ function Base.getindex(d::DArray, i::Int)
632652 _scalarindexingallowed ()
633653 return getindex_tuple (d, Tuple (CartesianIndices (d)[i]))
634654end
635- function Base. getindex (d:: DArray , i:: Int... )
655+ function Base. getindex (d:: DArray{<:Any,N} , i:: Vararg{ Int,N} ) where {N}
636656 _scalarindexingallowed ()
637657 return getindex_tuple (d, i)
638658end
639-
640659Base. getindex (d:: DArray ) = d[1 ]
641660Base. getindex (d:: SubDArray , I:: Int... ) = invoke (getindex, Tuple{SubArray{<: Any ,N},Vararg{Int,N}} where N, d, I... )
642661Base. getindex (d:: SubOrDArray , I:: Union{Int,UnitRange{Int},Colon,Vector{Int},StepRange{Int,Int}} ...) = view (d, I... )
0 commit comments