Merge pull request #250 from JuliaGPU/tb/map_args

Support GPU-like array types with map.

Author: maleadt

src/host/broadcast.jl

@@ -81,12 +81,12 @@ end
 allequal(x) = true
 allequal(x, y, z...) = x == y && allequal(y, z...)
 
-function Base.map!(f, y::AbstractGPUArray, xs::AbstractArray...)
+function Base.map!(f, y::GPUDestArray, xs::AbstractArray...)
     @assert allequal(size.((y, xs...))...)
     return y .= f.(xs...)
 end
 
-function Base.map(f, y::AbstractGPUArray, xs::AbstractArray...)
+function Base.map(f, y::GPUDestArray, xs::AbstractArray...)
     @assert allequal(size.((y, xs...))...)
     return f.(y, xs...)
 end

src/reference.jl

@@ -212,9 +212,9 @@ JLArrayStyle{M}(::Val{N}) where {N,M} = JLArrayStyle{N}()
 
 BroadcastStyle(::Type{JLArray{T,N}}) where {T,N} = JLArrayStyle{N}()
 
+# Allocating the output container
 Base.similar(bc::Broadcasted{JLArrayStyle{N}}, ::Type{T}) where {N,T} =
     similar(JLArray{T}, axes(bc))
-
 Base.similar(bc::Broadcasted{JLArrayStyle{N}}, ::Type{T}, dims...) where {N,T} =
     JLArray{T}(undef, dims...)