Customizing Broadcasting Recursively for Collections of Arrays

FYI, in terms of automatically broadcasting things on the right hand side, it’s kind of easy to get automatic broadcasting for arbitrary (“well-behaving”) struct. To do this, you can just subtype the following BroadcastableStruct (which is just a subset of BroadcastableCallable I described here):

abstract type BroadcastableStruct end

fieldvalues(obj) = ntuple(i -> getfield(obj, i), fieldcount(typeof(obj)))

# Taken from `Setfield.constructor_of`:
@generated constructor_of(::Type{T}) where T =
    getfield(parentmodule(T), nameof(T))

Broadcast.broadcastable(obj::BroadcastableStruct) =
    Broadcast.broadcasted(
        constructor_of(typeof(obj)),
        map(Broadcast.broadcastable, fieldvalues(obj))...)

Usage:

struct A{T} <: BroadcastableStruct
    a::T
    b::T
end

f(x, y) = A(x.a * y.a, x.b * y.b)
g(x::A) = x.a + x.b

g.(f.(A(ones(3), ones(3)), A(ones(3), ones(3))))

Result:

3-element Array{Float64,1}:
 2.0
 2.0
 2.0

Of course, you need to make sure that the construction of the struct is “free” (or define such one and use it via constructor_of). It’s just a convenient way to use broadcasting so I guess it does not improve any performance, though (I haven’t checked how it works with the inference etc).

I don’t think it’s hard to define copyto!(::A{<:AbstractArray{T}}, bc::Broadcasted) for the case “eltype” of bc is A{T}. But this is a bit more tricky to do at the level of BroadcastableStruct.

2 Likes