Broadcast for vectors of vector like objects

Say I have an object Foo for which broadcast is defined

struct Foo{T, V<:AbstractVector{T}} 
    data::V
end

Foo(data::AbstractVector) = Foo{eltype(data), typeof(data)}(data)

@generated function Base.Broadcast.broadcast!(f, u::Foo, args::Vararg{Any, n}) where {n}
    quote 
        $(Expr(:meta, :inline))
        broadcast!(f, unsafe_get(u), map(unsafe_get, args)...)
        u
    end
end

Base.unsafe_get(u::Foo) = u.data

Say I have two vectors of Foos, as in

w = [Foo([1, 2]), Foo([3, 4])]
v = [Foo([2, 3]), Foo([5, 6])]

I would like to have code such that this broadcast operation

w .+= v

works, and gives

w = [Foo([3, 5]), Foo([8, 10])]

with no temporary allocations, because the dot notation is transferred efficiently down a level to Foo objects. A quick look on github shows that there is an open issue about this, that means it is a known thing.

Has anyone found a good way to express this? Rather than a Vector of Foos, in actual code I have a custom container type wrapping a Vector so I can optionally define broadcast for this wrapper as well if that is needed.

I use a ArrayPartitions for this kind of thing all the time:

VectorOfArrays will also do it, but it will use the indexing because the broadcast PR isn’t done yet. But broadcasting that kind of thing would use the same overloads as an ArrayPartition, so that code is somewhat a tutorial right now.

1 Like