Broadcast for custom type

I am using v0.6, this is an MWE of a question I want to solve.

I would like to make broadcast work for a custom type

struct CongruentVector{Td,Tv}
    descriptor::Td
    vector::Tv
end

the following way: given a method with fallback

≅(a, b) = a == b

I would test that all CongruentVector arguments to broadcast have descriptors that are , then broadcast on their vectors and nothing else, ie every other argument would be treated as a “scalar”.

The idea behind this interface is that it would enforce matching descriptors. In particular, I would use it for posterior analysis on MCMC, and want to ensure that vectors (which are posterior draws) really do come from the same chains(s). Compatible dimensions are necessary but not sufficient for this.

Here is some mock code (that is incomplete as is):

function _common_descriptor(y, xs...)
    @assert all(y ≅ x for x in xs) "Arguments are not ≅."
    y
end

@generated function broadcast(f, xs::FIXMEunsurehowtotype...)
    ex_descriptors = tuple((:(xs[$i]).descriptor for (i,x) in enumerate(xs)
                            if x <: CongruentVector)...)
    ex_arguments = tuple((x <: CongruentVector ? :(xs[$i].vector) : :(Ref(xs[$i]))
                          for (i,x) in enumerate(xs))...)
    quote
        descriptor = common_descriptor($ex_descriptors...)
        result = broadcast(f, $ex_arguments...)
        CongruentVector(descriptor, result)
    end
end

From discussions I learned that I should probably be using generated functions. But it is unclear how to write the signature of a method that is only called when at least some arguments are CongruentVectors, but otherwise would not interfere with broadcast. Any pointers would be appreciated, including suggestions that I should use some other approach because what I am attempting is not currently feasible.

This issue came up recently for StaticArrays.jl as well. I wasn’t involved in figuring it out, but I remember the discussion here: https://github.com/JuliaArrays/StaticArrays.jl/pull/136#discussion_r109543567

1 Like

This does need to be documented better. In the meantime, consider the following pointers that may help:

This is a lightly edited version of a “minimal working example” I posted on Gitter, during a discussion with @ChrisRackauckas :

Depending on your use case you can choose to use as much or as little of the generic broadcast code as is necessary. Generally, you do not specialize broadcast, but rather broadcast_c.

I have a quick demo of what broadcast_c means:

julia> struct Poison end

julia> Base.Broadcast._containertype(::Type{<:Poison}) = Poison

julia> Base.Broadcast.promote_containertype(::Type{Poison}, _) = Poison

julia> Base.Broadcast.promote_containertype(_, ::Type{Poison}) = Poison

julia> Base.Broadcast.promote_containertype(::Type{Poison}, ::Type{Array}) = Poison

julia> Base.Broadcast.promote_containertype(::Type{Array}, ::Type{Poison}) = Poison

julia> Base.Broadcast.broadcast_c(f, ::Type{Poison}, _...) = "hijacked broadcasting"

julia> Poison() .+ [1, 2, 3]
"hijacked broadcasting"
  • roughly, the way broadcast works is: it looks at the argument types it gets and tries to determine which broadcast_c method handles broadcasting of those types.
  • first it calls _containertype on the types of all arguments
  • then it calls promote_containertype on all of the return values of _containertype

So the part to ensure that Broadcast will dispatch to your desired broadcast_c method is:

struct Poison end
Base.Broadcast._containertype(::Type{<:Poison}) = Poison
Base.Broadcast.promote_containertype(::Type{Poison}, _) = Poison
Base.Broadcast.promote_containertype(_, ::Type{Poison}) = Poison

and the part actually defining the broadcast behavior is:

Base.Broadcast.broadcast_c(f, ::Type{Poison}, _...) = "hijacked broadcasting"

Why do we need to treat arrays separately? In this case, because of ambiguities. The broadcast code in Base itself has the equivalent of

promote_containertype(::Type{Array}, _) = Array
promote_containertype(_, ::Type{Array}) = Array

so there will be ambiguities unless you treat arrays with their own extra methods.

Note broadcast dispatches to broadcast_c and broadcast!(f, ::AbstractArray, ...) dispatches to broadcast_c! in roughly the same way. Note that broadcast! by default is only defined for the LHS an AbstractArray. If you have your own type on the LHS you will need to specialize broadcast! directly instead of broadcast_c! (or perhaps in addition to, if you’d like to call broadcast_c!).

So in summary:

  • Do not add methods to broadcast.
  • Instead, add methods to broadcast_c, which has the same signature of broadcast except an additional argument is added after the function argument, which will be a ::Type{...} of your choosing.
  • Then, add methods to _containertype and promote_containertype so that the generic broadcast can pick the correct method.
  • The same applies for broadcast! and broadcast_c!, except that when the LHS is not an AbstractArray, additional methods may be required.
16 Likes

@fengyang.wang That description is really great. Would you add it to the developer docs in the manual?

10 Likes

Seems like a good idea. I’ll fix it up over the weekend and make a PR.

In the link you provided to DataValues.jl, but also in Julia Base broadcast.jl, there are also broadcast_indices(...), and _broadcast_getindex_eltype(...) etc. Are these intended as potential customization points for indices and element types when broadcasting?

Those are used so that DataValues behave like scalars with broadcast with arrays (the array broadcast_c calls the functions broadcast_indices and _broadcast_getindex_eltype(...)). They are needed if a combination of your type and arrays will promote_containertype to Array, but not otherwise.

Effectively, each broadcast_c method has its own points of potential/required customization for other types. When broadcasting your types with others, and deferring to the others’ broadcast_c method, you must be mindful of the interface the other broadcast_c method expects. This does result in a possible combinatorial explosion of complexity, which is something that is being worked on for 0.7 and 1.0.

1 Like

Would be great if you could also explain how to cope with cases where the LHS is not an AbstractArray! Thanks

‎In such cases you directly add methods to broadcast!, possibly reusing the generic infrastructure.

1 Like

First, thank you very much for this!

I did run into one issue. I think

Base.Broadcast.promote_containertype(::Type{Poison}, ::Type{Poison}) = Poison

should be added, because otherwise Base.Broadcast.promote_containertype(Poison, Poison) results in ambiguity. This happens for example when broadcast(+, Poison(), Poison()) is called.

2 Likes

Good old ambiguities! Thanks, I’ll add a note. Actually, that post is too old to edit. It should probably be replaced with proper documentation, in any event.

2 Likes