Skip to content

Commit

Permalink
fix #4883, result type of broadcast for arbitrary functions
Browse files Browse the repository at this point in the history
  • Loading branch information
JeffBezanson committed Jun 28, 2016
1 parent 0dcb477 commit 9a9c37c
Showing 1 changed file with 74 additions and 1 deletion.
75 changes: 74 additions & 1 deletion base/broadcast.jl
Original file line number Diff line number Diff line change
Expand Up @@ -140,7 +140,80 @@ end
B
end

@inline broadcast(f, As...) = broadcast!(f, allocate_for(Array{promote_eltype_op(f, As...)}, As, broadcast_shape(As...)), As...)
# broadcast with computed element type

@generated function _broadcast!{M,AT,nargs}(f, B::AbstractArray, indexmaps::M, As::AT, ::Type{Val{nargs}}, iter, st, count)
quote
$(Expr(:meta, :noinline))
# destructure the indexmaps and As tuples
@nexprs $nargs i->(A_i = As[i])
@nexprs $nargs i->(imap_i = indexmaps[i])
while !done(iter, st)
I, st = next(iter, st)
# reverse-broadcast the indices
@nexprs $nargs i->(I_i = newindex(I, imap_i))
# extract array values
@nexprs $nargs i->(@inbounds val_i = A_i[I_i])
# call the function
V = @ncall $nargs f val
S = typeof(V)
# store the result
if S <: eltype(B)
@inbounds B[I] = V
else
R = typejoin(eltype(B), S)
new = similar(B, R)
for II in take(iter, count)
new[II] = B[II]
end
new[I] = V
return _broadcast!(f, new, indexmaps, As, Val{nargs}, iter, st, count+1)
end
count += 1
end
return B
end
end

function broadcast_t(f, ::Type{Any}, As...)
shp = broadcast_shape(As...)
iter = CartesianRange(shp)
if isempty(iter)
return allocate_for(Array{Union{}}, As, shp)
end
nargs = length(As)
sz = size(iter)
indexmaps = map(x->newindexer(sz, x), As)
st = start(iter)
I, st = next(iter, st)
val = f([ As[i][newindex(I, indexmaps[i])] for i=1:nargs ]...)
B = allocate_for(Array{typeof(val)}, As, shp)
B[I] = val
return _broadcast!(f, B, indexmaps, As, Val{nargs}, iter, st, 1)
end

@inline broadcast_t(f, T, As...) = broadcast!(f, allocate_for(Array{T}, As, broadcast_shape(As...)), As...)

@inline broadcast(f, As...) = broadcast_t(f, promote_eltype_op(f, As...), As...)

# alternate, more compact implementation; unfortunately slower.
# also the `collect` machinery doesn't yet support arbitrary index bases.
#=
@generated function _broadcast{nargs}(f, indexmaps, As, ::Type{Val{nargs}}, iter)
quote
collect((@ncall $nargs f i->As[i][newindex(I, indexmaps[i])]) for I in iter)
end
end
function broadcast(f, As...)
shp = broadcast_shape(As...)
iter = CartesianRange(shp)
sz = size(iter)
indexmaps = map(x->newindexer(sz, x), As)
naT = Val{nfields(As)}
_broadcast(f, indexmaps, As, naT, iter)
end
=#

@inline bitbroadcast(f, As...) = broadcast!(f, allocate_for(BitArray, As, broadcast_shape(As...)), As...)

Expand Down

2 comments on commit 9a9c37c

@stevengj
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No tests?

@JeffBezanson
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is an earlier version; 4a9f6d9 has some tests. We're also still waiting on some updates to promote_op.

Please sign in to comment.