Skip to content

Commit

Permalink
Extend sparse broadcast! to combinations of broadcast scalars and spa…
Browse files Browse the repository at this point in the history
…rse vectors/matrices.

Makes broadcast! dispatch on container type (as broadcast), and inject generic sparse broadcast! for the appropriate container type.
  • Loading branch information
Sacha0 committed Dec 26, 2016
1 parent 18df7c5 commit d31145b
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 4 deletions.
3 changes: 2 additions & 1 deletion base/broadcast.jl
Original file line number Diff line number Diff line change
Expand Up @@ -190,7 +190,8 @@ Note that `dest` is only used to store the result, and does not supply
arguments to `f` unless it is also listed in the `As`,
as in `broadcast!(f, A, A, B)` to perform `A[:] = broadcast(f, A, B)`.
"""
@inline function broadcast!{nargs}(f, B::AbstractArray, As::Vararg{Any,nargs})
@inline broadcast!(f, B, As...) = broadcast!_c(f, containertype(B, As...), B, As...)
@inline function broadcast!_c{nargs}(f, ::Type, B::AbstractArray, As::Vararg{Any,nargs})
shape = indices(B)
check_broadcast_indices(shape, As...)
keeps, Idefaults = map_newindexer(shape, As)
Expand Down
11 changes: 8 additions & 3 deletions base/sparse/higherorderfns.jl
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,8 @@ module HigherOrderFns
# This module provides higher order functions specialized for sparse arrays,
# particularly map[!]/broadcast[!] for SparseVectors and SparseMatrixCSCs at present.
import Base: map, map!, broadcast, broadcast!
import Base.Broadcast: containertype, promote_containertype, broadcast_indices, broadcast_c
import Base.Broadcast: containertype, promote_containertype,
broadcast_indices, broadcast_c, broadcast!_c

using Base: front, tail, to_shape
using ..SparseArrays: SparseVector, SparseMatrixCSC, AbstractSparseArray, indtype
Expand Down Expand Up @@ -850,11 +851,15 @@ promote_containertype(::Type{Tuple}, ::Type{AbstractSparseArray}) = Array
promote_containertype(::Type{AbstractSparseArray}, ::Type{Array}) = Array
promote_containertype(::Type{AbstractSparseArray}, ::Type{Tuple}) = Array

# broadcast entry point for combinations of sparse arrays and other types
function broadcast_c(f, ::Type{AbstractSparseArray}, mixedargs...)
# broadcast[!] entry points for combinations of sparse arrays and other types
@inline function broadcast_c{N}(f, ::Type{AbstractSparseArray}, mixedargs::Vararg{Any,N})
parevalf, passedargstup = capturescalars(f, mixedargs)
return broadcast(parevalf, passedargstup...)
end
@inline function broadcast!_c{N}(f, ::Type{AbstractSparseArray}, dest::SparseVecOrMat, mixedsrcargs::Vararg{Any,N})
parevalf, passedsrcargstup = capturescalars(f, mixedsrcargs)
return broadcast!(parevalf, dest, passedsrcargstup...)
end
# capturescalars takes a function (f) and a tuple of mixed sparse vectors/matrices and
# broadcast scalar arguments (mixedargs), and returns a function (parevalf) and a reduced
# argument tuple (passedargstup) containing only the sparse vectors/matrices in mixedargs
Expand Down

0 comments on commit d31145b

Please sign in to comment.