Skip to content

Commit 6ebf703

Browse files
committed
Fix JuliaLang#24914 (WIP).
1 parent 9b8e651 commit 6ebf703

File tree

2 files changed

+29
-25
lines changed

2 files changed

+29
-25
lines changed

base/broadcast.jl

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -242,12 +242,6 @@ broadcast_indices
242242
# special cases defined for performance
243243
broadcast(f, x::Number...) = f(x...)
244244
@inline broadcast(f, t::NTuple{N,Any}, ts::Vararg{NTuple{N,Any}}) where {N} = map(f, t, ts...)
245-
@inline broadcast!(::typeof(identity), x::AbstractArray{T,N}, y::AbstractArray{S,N}) where {T,S,N} =
246-
Base.axes(x) == Base.axes(y) ? copy!(x, y) : _broadcast!(identity, x, y)
247-
248-
# special cases for "X .= ..." (broadcast!) assignments
249-
broadcast!(::typeof(identity), X::AbstractArray, x::Number) = fill!(X, x)
250-
broadcast!(f, X::AbstractArray, x::Number...) = (@inbounds for I in eachindex(X); X[I] = f(x...); end; X)
251245

252246
## logic for deciding the BroadcastStyle
253247
# Dimensionality: computing max(M,N) in the type domain so we preserve inferrability
@@ -448,8 +442,18 @@ Note that `dest` is only used to store the result, and does not supply
448442
arguments to `f` unless it is also listed in the `As`,
449443
as in `broadcast!(f, A, A, B)` to perform `A[:] = broadcast(f, A, B)`.
450444
"""
451-
@inline broadcast!(f, C::AbstractArray, A, Bs::Vararg{Any,N}) where {N} =
452-
_broadcast!(f, C, A, Bs...)
445+
broadcast!(f, dest, As...) = broadcast!(f, dest, combine_styles(As...), As...)
446+
broadcast!(f, dest, ::BroadcastStyle, As...) = broadcast!(f, dest, nothing, As...)
447+
@inline function broadcast!(f, C, ::Void, A, Bs::Vararg{Any,N}) where N
448+
if isa(f, typeof(identity)) && N == 0
449+
if isa(A, Number)
450+
return fill!(C, A)
451+
elseif isa(C, AbstractArray) && isa(A, AbstractArray) && Base.axes(C) == Base.axes(A)
452+
return copy!(C, A)
453+
end
454+
end
455+
return _broadcast!(f, C, A, Bs...)
456+
end
453457

454458
# This indirection allows size-dependent implementations (e.g., see the copying `identity`
455459
# specialization above)

base/sparse/higherorderfns.jl

Lines changed: 17 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,8 @@ end
9393
# (3) broadcast[!] entry points
9494
broadcast(f::Tf, A::SparseVector) where {Tf} = _noshapecheck_map(f, A)
9595
broadcast(f::Tf, A::SparseMatrixCSC) where {Tf} = _noshapecheck_map(f, A)
96-
function broadcast!(f::Tf, C::SparseVecOrMat) where Tf
96+
97+
function broadcast!(f::Tf, C::SparseVecOrMat, ::Void) where Tf
9798
isempty(C) && return _finishempty!(C)
9899
fofnoargs = f()
99100
if _iszero(fofnoargs) # f() is zero, so empty C
@@ -106,14 +107,13 @@ function broadcast!(f::Tf, C::SparseVecOrMat) where Tf
106107
end
107108
return C
108109
end
109-
function broadcast!(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat, Bs::Vararg{SparseVecOrMat,N}) where {Tf,N}
110-
_aresameshape(C, A, Bs...) && return _noshapecheck_map!(f, C, A, Bs...)
111-
Base.Broadcast.check_broadcast_indices(axes(C), A, Bs...)
112-
fofzeros = f(_zeros_eltypes(A, Bs...)...)
113-
fpreszeros = _iszero(fofzeros)
114-
return fpreszeros ? _broadcast_zeropres!(f, C, A, Bs...) :
115-
_broadcast_notzeropres!(f, fofzeros, C, A, Bs...)
110+
function broadcast!(f, dest::SparseVecOrMat, ::Void, A, Bs::Vararg{Any,N}) where N
111+
if isa(f, typeof(identity)) && N == 0 && isa(A, Number)
112+
return fill!(dest, A)
113+
end
114+
return spbroadcast_args!(f, dest, Broadcast.combine_styles(A, Bs...), A, Bs...)
116115
end
116+
117117
# the following three similar defs are necessary for type stability in the mixed vector/matrix case
118118
broadcast(f::Tf, A::SparseVector, Bs::Vararg{SparseVector,N}) where {Tf,N} =
119119
_aresameshape(A, Bs...) ? _noshapecheck_map(f, A, Bs...) : _diffshape_broadcast(f, A, Bs...)
@@ -1005,18 +1005,18 @@ Broadcast.BroadcastStyle(::SparseMatStyle, ::Broadcast.DefaultArrayStyle{N}) whe
10051005
broadcast(f, ::PromoteToSparse, ::Void, ::Void, As::Vararg{Any,N}) where {N} =
10061006
broadcast(f, map(_sparsifystructured, As)...)
10071007

1008-
# ambiguity resolution
1009-
broadcast!(::typeof(identity), dest::SparseVecOrMat, x::Number) =
1010-
fill!(dest, x)
1011-
broadcast!(f, dest::SparseVecOrMat, x::Number...) =
1012-
spbroadcast_args!(f, dest, SPVM, x...)
1013-
10141008
# For broadcast! with ::Any inputs, we need a layer of indirection to determine whether
10151009
# the inputs can be promoted to SparseVecOrMat. If it's just SparseVecOrMat and scalars,
10161010
# we can handle it here, otherwise see below for the promotion machinery.
1017-
broadcast!(f, dest::SparseVecOrMat, mixedsrcargs::Vararg{Any,N}) where N =
1018-
spbroadcast_args!(f, dest, Broadcast.combine_styles(mixedsrcargs...), mixedsrcargs...)
1019-
function spbroadcast_args!(f, dest, ::Type{SPVM}, mixedsrcargs::Vararg{Any,N}) where N
1011+
function spbroadcast_args!(f::Tf, C, ::SPVM, A::SparseVecOrMat, Bs::Vararg{SparseVecOrMat,N}) where {Tf,N}
1012+
_aresameshape(C, A, Bs...) && return _noshapecheck_map!(f, C, A, Bs...)
1013+
Base.Broadcast.check_broadcast_indices(axes(C), A, Bs...)
1014+
fofzeros = f(_zeros_eltypes(A, Bs...)...)
1015+
fpreszeros = _iszero(fofzeros)
1016+
return fpreszeros ? _broadcast_zeropres!(f, C, A, Bs...) :
1017+
_broadcast_notzeropres!(f, fofzeros, C, A, Bs...)
1018+
end
1019+
function spbroadcast_args!(f, dest, ::SPVM, mixedsrcargs::Vararg{Any,N}) where N
10201020
# mixedsrcargs contains nothing but SparseVecOrMat and scalars
10211021
parevalf, passedsrcargstup = capturescalars(f, mixedsrcargs)
10221022
return broadcast!(parevalf, dest, passedsrcargstup...)

0 commit comments

Comments
 (0)