Skip to content

Commit 746dbb0

Browse files
committed
Extend sparse broadcast! to combinations of broadcast scalars and sparse vectors/matrices.
Makes broadcast! dispatch on container type (as broadcast), and inject generic sparse broadcast! for the appropriate container type.
1 parent b9cc65b commit 746dbb0

File tree

2 files changed

+11
-4
lines changed

2 files changed

+11
-4
lines changed

base/broadcast.jl

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -202,7 +202,9 @@ Note that `dest` is only used to store the result, and does not supply
202202
arguments to `f` unless it is also listed in the `As`,
203203
as in `broadcast!(f, A, A, B)` to perform `A[:] = broadcast(f, A, B)`.
204204
"""
205-
@inline function broadcast!{N}(f, C::AbstractArray, A, Bs::Vararg{Any,N})
205+
@inline broadcast!{N}(f, C::AbstractArray, A, Bs::Vararg{Any,N}) =
206+
broadcast_c!(f, containertype(C, A, Bs...), C, A, Bs...)
207+
@inline function broadcast_c!{N}(f, ::Type, C::AbstractArray, A, Bs::Vararg{Any,N})
206208
shape = indices(C)
207209
@boundscheck check_broadcast_indices(shape, A, Bs...)
208210
keeps, Idefaults = map_newindexer(shape, A, Bs)

base/sparse/higherorderfns.jl

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,8 @@ module HigherOrderFns
55
# This module provides higher order functions specialized for sparse arrays,
66
# particularly map[!]/broadcast[!] for SparseVectors and SparseMatrixCSCs at present.
77
import Base: map, map!, broadcast, broadcast!
8-
import Base.Broadcast: containertype, promote_containertype, broadcast_indices, broadcast_c
8+
import Base.Broadcast: containertype, promote_containertype,
9+
broadcast_indices, broadcast_c, broadcast_c!
910

1011
using Base: front, tail, to_shape
1112
using ..SparseArrays: SparseVector, SparseMatrixCSC, AbstractSparseArray, indtype
@@ -852,11 +853,15 @@ promote_containertype(::Type{Tuple}, ::Type{AbstractSparseArray}) = Array
852853
promote_containertype(::Type{AbstractSparseArray}, ::Type{Array}) = Array
853854
promote_containertype(::Type{AbstractSparseArray}, ::Type{Tuple}) = Array
854855

855-
# broadcast entry point for combinations of sparse arrays and other types
856-
function broadcast_c(f, ::Type{AbstractSparseArray}, mixedargs...)
856+
# broadcast[!] entry points for combinations of sparse arrays and other types
857+
@inline function broadcast_c{N}(f, ::Type{AbstractSparseArray}, mixedargs::Vararg{Any,N})
857858
parevalf, passedargstup = capturescalars(f, mixedargs)
858859
return broadcast(parevalf, passedargstup...)
859860
end
861+
@inline function broadcast_c!{N}(f, ::Type{AbstractSparseArray}, dest::SparseVecOrMat, mixedsrcargs::Vararg{Any,N})
862+
parevalf, passedsrcargstup = capturescalars(f, mixedsrcargs)
863+
return broadcast!(parevalf, dest, passedsrcargstup...)
864+
end
860865
# capturescalars takes a function (f) and a tuple of mixed sparse vectors/matrices and
861866
# broadcast scalar arguments (mixedargs), and returns a function (parevalf) and a reduced
862867
# argument tuple (passedargstup) containing only the sparse vectors/matrices in mixedargs

0 commit comments

Comments
 (0)