Skip to content

Commit

Permalink
Use broadcast for element wise operators where appropriate
Browse files Browse the repository at this point in the history
  • Loading branch information
pabloferz committed Sep 28, 2016
1 parent 951aea0 commit 20f5ba0
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 8 deletions.
15 changes: 8 additions & 7 deletions base/broadcast.jl
Original file line number Diff line number Diff line change
Expand Up @@ -240,7 +240,7 @@ end
end

# broadcast methods that dispatch on the type found by inference
function _broadcast_t(f, ::Type{Any}, shape, iter, As...)
function broadcast_t(f, ::Type{Any}, shape, iter, As...)
nargs = length(As)
keeps, Idefaults = map_newindexer(shape, As)
st = start(iter)
Expand All @@ -250,7 +250,7 @@ function _broadcast_t(f, ::Type{Any}, shape, iter, As...)
B[I] = val
return _broadcast!(f, B, keeps, Idefaults, As, Val{nargs}, iter, st, 1)
end
@inline function _broadcast_t(f, T, shape, iter, As...)
@inline function broadcast_t(f, T, shape, iter, As...)
B = similar(Array{T}, shape)
nargs = length(As)
keeps, Idefaults = map_newindexer(shape, As)
Expand All @@ -260,7 +260,8 @@ end

# broadcast method that uses inference to find the type, but preserves abstract
# container types when possible (used by binary elementwise operators)
@inline broadcast_t(f, As...) = broadcast!(f, similar(Array{promote_eltype_op(f, As...)}, broadcast_indices(As...)), As...)
@inline broadcast_elwise_op(f, As...) =
broadcast!(f, similar(Array{promote_eltype_op(f, As...)}, broadcast_indices(As...)), As...)

ftype(f, A) = typeof(a->f(a))
ftype(f, A...) = typeof(a->f(a...))
Expand All @@ -272,7 +273,7 @@ ziptype(A, B) = Zip2{Tuple{eltype(A)},Tuple{eltype(B)}}

# broadcast methods that dispatch on the type of the final container
@inline function broadcast_c(f, ::Type{Array}, As...)
S = _default_eltype(Generator{ziptype(As...),ftype(f, As...)})
T = _default_eltype(Generator{ziptype(As...),ftype(f, As...)})
shape = broadcast_indices(As...)
iter = CartesianRange(shape)
if isleaftype(T)
Expand Down Expand Up @@ -471,10 +472,10 @@ end
## elementwise operators ##

for op in (:÷, :%, :<<, :>>, :-, :/, :\, ://, :^)
@eval $(Symbol(:., op))(A::AbstractArray, B::AbstractArray) = broadcast_t($op, A, B)
@eval $(Symbol(:., op))(A::AbstractArray, B::AbstractArray) = broadcast_elwise_op($op, A, B)
end
.+(As::AbstractArray...) = broadcast_t(+, As...)
.*(As::AbstractArray...) = broadcast_t(*, As...)
.+(As::AbstractArray...) = broadcast_elwise_op(+, As...)
.*(As::AbstractArray...) = broadcast_elwise_op(*, As...)

# ## element-wise comparison operators returning BitArray ##

Expand Down
2 changes: 1 addition & 1 deletion base/sparse/sparsematrix.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1746,7 +1746,7 @@ broadcast{Tv1,Ti1,Tv2,Ti2}(f::Function, A_1::SparseMatrixCSC{Tv1,Ti1}, A_2::Spar
broadcast!(f, spzeros(promote_type(Tv1, Tv2), promote_type(Ti1, Ti2), to_shape(broadcast_indices(A_1, A_2))), A_1, A_2)

@inline broadcast_zpreserving!(args...) = broadcast!(args...)
@inline broadcast_zpreserving(args...) = broadcast(args...)
@inline broadcast_zpreserving(args...) = Base.Broadcast.broadcast_elwise_op(args...)
broadcast_zpreserving{Tv1,Ti1,Tv2,Ti2}(f::Function, A_1::SparseMatrixCSC{Tv1,Ti1}, A_2::SparseMatrixCSC{Tv2,Ti2}) =
broadcast_zpreserving!(f, spzeros(promote_type(Tv1, Tv2), promote_type(Ti1, Ti2), to_shape(broadcast_indices(A_1, A_2))), A_1, A_2)
broadcast_zpreserving{Tv,Ti}(f::Function, A_1::SparseMatrixCSC{Tv,Ti}, A_2::Union{Array,BitArray,Number}) =
Expand Down

0 comments on commit 20f5ba0

Please sign in to comment.