Skip to content

Commit

Permalink
Use broadcast for array ops, closes #17254
Browse files Browse the repository at this point in the history
  • Loading branch information
martinholters committed Jul 11, 2016
1 parent efa23ad commit b7c1bd9
Show file tree
Hide file tree
Showing 3 changed files with 37 additions and 64 deletions.
42 changes: 13 additions & 29 deletions base/arraymath.jl
Original file line number Diff line number Diff line change
Expand Up @@ -50,37 +50,21 @@ end
@pure promote_array_type{S<:Integer, P}(F, ::Type{S}, ::Type{Bool}, ::Type{P}) = P

for f in (:+, :-, :div, :mod, :&, :|, :$)
@eval begin
function ($f){S,T}(A::Range{S}, B::Range{T})
F = similar(A, promote_op($f,S,T), promote_shape(size(A),size(B)))
for (iF, iA, iB) in zip(eachindex(F), eachindex(A), eachindex(B))
@inbounds F[iF] = ($f)(A[iA], B[iB])
end
return F
end
function ($f){S,T}(A::AbstractArray{S}, B::Range{T})
F = similar(A, promote_op($f,S,T), promote_shape(A,B))
for (iF, iA, iB) in zip(eachindex(F), eachindex(A), eachindex(B))
@inbounds F[iF] = ($f)(A[iA], B[iB])
end
return F
end
function ($f){S,T}(A::Range{S}, B::AbstractArray{T})
F = similar(B, promote_op($f,S,T), promote_shape(A,B))
for (iF, iA, iB) in zip(eachindex(F), eachindex(A), eachindex(B))
@inbounds F[iF] = ($f)(A[iA], B[iB])
end
return F
end
function ($f){S,T}(A::AbstractArray{S}, B::AbstractArray{T})
F = similar(A, promote_op($f,S,T), promote_shape(A,B))
for (iF, iA, iB) in zip(eachindex(F), eachindex(A), eachindex(B))
@inbounds F[iF] = ($f)(A[iA], B[iB])
end
return F
end
@eval ($f){S,T}(A::AbstractArray{S}, B::AbstractArray{T}) =
_elementwise($f, A, B, promote_eltype_op($f, A, B))
end
function _elementwise{S,T}(op, A::AbstractArray{S}, B::AbstractArray{T}, ::Type{Any})
promote_shape(A,B) # check size compatibility
return broadcast(op, A, B)
end
function _elementwise{S,T,R}(op, A::AbstractArray{S}, B::AbstractArray{T}, ::Type{R})
F = similar(A, R, promote_shape(A,B))
for (iF, iA, iB) in zip(eachindex(F), eachindex(A), eachindex(B))
@inbounds F[iF] = op(A[iA], B[iB])
end
return F
end

for f in (:.+, :.-, :.*, :./, :.\, :.^, :, :.%, :.<<, :.>>, :div, :mod, :rem, :&, :|, :$)
@eval begin
function ($f){T}(A::Number, B::AbstractArray{T})
Expand Down
40 changes: 5 additions & 35 deletions base/broadcast.jl
Original file line number Diff line number Diff line change
Expand Up @@ -266,44 +266,14 @@ end

## elementwise operators ##

(A::AbstractArray, B::AbstractArray) = broadcast(÷, A, B)
.%(A::AbstractArray, B::AbstractArray) = broadcast(%, A, B)
.<<(A::AbstractArray, B::AbstractArray) = broadcast(<<, A, B)
.>>(A::AbstractArray, B::AbstractArray) = broadcast(>>, A, B)

# should this be deprecated? Only use in Base is in sparsematrix.jl
eltype_plus(As::AbstractArray...) = promote_eltype_op(+, As...)

.+(As::AbstractArray...) = broadcast!(+, Array{eltype_plus(As...)}(to_shape(broadcast_shape(As...))), As...)

function .-(A::AbstractArray, B::AbstractArray)
broadcast!(-, Array{promote_op(-, eltype(A), eltype(B))}(to_shape(broadcast_shape(A,B))), A, B)
end

eltype_mul(As::AbstractArray...) = promote_eltype_op(*, As...)

.*(As::AbstractArray...) = broadcast!(*, Array{eltype_mul(As...)}(to_shape(broadcast_shape(As...))), As...)

function ./(A::AbstractArray, B::AbstractArray)
broadcast!(/, Array{promote_op(/, eltype(A), eltype(B))}(to_shape(broadcast_shape(A, B))), A, B)
end

function .\(A::AbstractArray, B::AbstractArray)
broadcast!(\, Array{promote_op(\, eltype(A), eltype(B))}(to_shape(broadcast_shape(A, B))), A, B)
end

typealias RatIntT{T<:Integer} Union{Type{Rational{T}},Type{T}}
typealias CRatIntT{T<:Integer} Union{Type{Complex{Rational{T}}},Type{Complex{T}},Type{Rational{T}},Type{T}}
type_rdiv{T<:Integer,S<:Integer}(::RatIntT{T}, ::RatIntT{S}) =
Rational{promote_type(T,S)}
type_rdiv{T<:Integer,S<:Integer}(::CRatIntT{T}, ::CRatIntT{S}) =
Complex{Rational{promote_type(T,S)}}
function .//(A::AbstractArray, B::AbstractArray)
broadcast!(//, Array{type_rdiv(eltype(A), eltype(B))}(to_shape(broadcast_shape(A, B))), A, B)
end

function .^(A::AbstractArray, B::AbstractArray)
broadcast!(^, Array{promote_op(^, eltype(A), eltype(B))}(to_shape(broadcast_shape(A, B))), A, B)
for op in (:÷, :%, :<<, :>>, :-, :/, :\, ://, :^)
@eval $(Symbol(:., op))(A::AbstractArray, B::AbstractArray) = broadcast($(op), A, B)
end
.+(As::AbstractArray...) = broadcast(+, As...)
.*(As::AbstractArray...) = broadcast(*, As...)

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

Expand Down
19 changes: 19 additions & 0 deletions test/arrayops.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1668,3 +1668,22 @@ let A = zeros(3,3)
@test size(A[:,UInt(1):UInt(2)]) == (3,2)
@test size(similar(A, UInt(3), 0x3)) == size(similar(A, (UInt(3), 0x3))) == (3,3)
end

# issue 17254
module AutoRetType

using Base.Test

immutable Foo end
for op in (:+, :*, :÷, :%, :<<, :>>, :-, :/, :\, ://, :^)
@eval import Base.$(op)
@eval $(op)(::Foo, ::Foo) = Foo()
end
A = fill(Foo(), 10, 10)
@test typeof(A+A) == Matrix{Foo}
@test typeof(A-A) == Matrix{Foo}
for op in (:.+, :.*, :, :.%, :.<<, :.>>, :.-, :./, :.\, :.//, :.^)
@eval @test typeof($(op)(A,A)) == Matrix{Foo}
end

end

0 comments on commit b7c1bd9

Please sign in to comment.