Asking for code review (extending broadcast)

Internals & Design
1

I thought I would make an attempt at #23734, but before submitting even a WIP PR, I wanted to ask for some preliminary help, because I am not sure I grasp all the details in Base.Broadcast.

Could the following be a viable approach? I am trying to build on existing infrastructure as much as possible.

"Helper structure, maps tuple elements to elements in a tuple of arrays."
struct TupleofArrays{S <: Tuple,T,N} <: AbstractArray{T,N}
    arrays::S
    function TupleofArrays(arrays::S) where S
        @assert !isempty(arrays)
        siz = size(first(arrays))
        ixs = indices(first(arrays))
        for a in Base.tail(arrays)
            @assert size(a) == siz "Incompatible dimensions"
            @assert indices(a) == ixs "Incompatible indexing"
        end
        T = Tuple{map(eltype, arrays)...}
        new{S,T,length(siz)}(arrays)
    end
end

Base.size(toa::TupleofArrays) = size(toa.arrays[1])
Base.getindex(toa::TupleofArrays, ixs...) = map(a->a[ixs...], toa.arrays)
Base.eltype(toa::TupleofArrays{S,T,N}) where {S,T,N} = T
Base.ndims(toa::TupleofArrays{S,T,N}) where {S,T,N} = N
Base.setindex!(toa::TupleofArrays, value, ixs...) =
    map((a,v)->setindex!(a, v, ixs...), toa.arrays, value)
Base.similar(::Type{TupleofArrays}, eltype, shape) =
    TupleofArrays(map(t->similar(Array{t}, shape), tuple(eltype.parameters...)))

function Base.broadcast!(f, result::NTuple{N,AbstractArray}, args...) where N
    toa = TupleofArrays(result)
    broadcast!(f, toa, args...)
    toa.arrays
end

function broadcast_mval(f, args...)
    T = Base.Broadcast._broadcast_eltype(f, args...)
    shape = Base.Broadcast.broadcast_indices(args...)
    toa = similar(TupleofArrays, T, shape)
    broadcast!(f, toa, args...)
    toa.arrays
end

# demo
f(x) = x+1, x+2
x = reshape(1:10, :, 2)

# broadcast to multiple values (tuple of arrays)
a, b = broadcast_mval(f, x)

# broadcast into existing arrays -- no new syntax
a2 = similar(x)
b2 = similar(x)
a2, b2 .= f.(x)
2

Worked on it a bit and put the code in a package:
https://github.com/tpapp/MultiBroadcast.jl/
It works for me in a satisfactory way, tests run on 0.6 and nightly, and I fixed some quirks (eg bits mapping to a BitArray). If someone familiar with the internals of Base.Broadcast could comment, that would be great.

I have introduced an experimental macro called @multi, which works like this:

f(x, y) = x+y, x*y
a, b = @multi f.(1:3, 4:6)

Currently the way it works is that it expands the expression, then replaces the outermost broadcast with multi_broadcast. Is this the right approach — is it OK if a macro calls expand?

2 Likes