Best practice for broadcastable callables

I wonder what is the best practice for callables that are broadcastable for input and its parameters. Here is an example:

struct Scale{T}
    a::T
end

struct Shift{T}
    a::T
end

(f::Scale)(x) = f.a * x
(f::Shift)(x) = f.a + x

struct VCompose{F, G}
    f::F
    g::G
end

(f::VCompose)(x) = f.f.(f.g.(x))

Scale and Shift can be broadcasted as usual when their parameters are scalars:

julia> f = VCompose(Scale(10), Shift(0.1))
       f([1, 2, 3])
3-element Array{Float64,1}:
 11.0
 21.0
 31.0

However, this does not work when the parameters are vectors

julia> f = VCompose(Scale([10, 100, 1000]), Shift([10, 20, 30]))
       f([1, 2, 3])
ERROR: MethodError: no method matching +(::Array{Int64,1}, ::Int64)

Note that using dot expressions inside the call definition like below is not the right solution.

(f::Scale)(x) = f.a .* x
(f::Shift)(x) = f.a .+ x

This is because it would “doubly” broadcast the input.

julia> f = VCompose(Scale([10, 100, 1000]), Shift([10, 20, 30]))
       f([1, 2, 3])
3-element Array{Array{Int64,1},1}:
 [110, 2100, 31000]
 [120, 2200, 32000]
 [130, 2300, 33000]

What is the best way to do it?

One thing I can do is to overload broadcasted. That is to say, I can define

Broadcast.broadcasted(f::Scale, x) = Broadcast.broadcasted(*, f.a, x)
Broadcast.broadcasted(f::Shift, x) = Broadcast.broadcasted(+, f.a, x)

so that

julia> f = VCompose(Scale([10, 100, 1000]), Shift([10, 20, 30]))
       f([1, 2, 3])
3-element Array{Int64,1}:
   110
  2200
 33000

I’m using Broadcast.broadcasted directly here and it is hard to read. But I’m not worried about this aspect because Julia may implement “lazy” broadcasting feature and there is a macro that does this. Also, LazyArrays.jl may implement it soon.

On the other hand, it’s a bit frustrating that I have to implement the callable twice. It’s a bad idea to keep them consistent by just being careful. So, it probably is better to define the normal callable in terms of broadcasted specialization:

(f::Scale)(x) = f.(x)
(f::Shift)(x) = f.(x)

It is a bit strange definition because you’d think it would have a stack overflow. Or, if you know Julia well, maybe you’d assume there are scalar specializations for those callables. It could be very annoying to read such code but I guess I can solve it by just leaving comments.

I think this way of writing broadcastable callables is somewhat OK. (It actually is fun.) But, I’d like to know if there are better ways to do this.

2 Likes

How about just using dot expressions for Scale and Shift, but not Compose?

struct Scale{T} a::T end

struct Shift{T} a::T end

(f::Scale)(x) = f.a .* x
(f::Shift)(x) = f.a .+ x

struct VCompose{F, G} f::F; g::G end

(f::VCompose)(x) = f.f(f.g(x))

It seems like that would support all combinations:

julia> f = VCompose(Scale(10), Shift(0.1));

julia> f(1)
11.0

julia> f = VCompose(Scale(10), Shift(0.1));

julia> f([1, 2, 3])
3-element Array{Float64,1}:
 11.0
 21.0
 31.0

julia> f = VCompose(Scale([10, 100, 1000]), Shift([10, 20, 30]));

julia> f(1)
3-element Array{Int64,1}:
   110
  2100
 31000

julia> f = VCompose(Scale([10, 100, 1000]), Shift([10, 20, 30]));

julia> f([1, 2, 3])
3-element Array{Int64,1}:
   110
  2200
 33000

Yes, good point. I forgot to mention that I want the callables Scale and Shift to be fusible to avoid allocations. I guess VCompose was a misleading example because it allocates always.

For example, I can define

vapply!(f::VCompose, y, x) = y .= f.f.(f.g.(x))

to avoid allocations in my example.

I think I’ve found a better approach. Implementing a scalar operation based on a broadcasted operation felt backward and sometimes were hard to work with. It turned out it’s rather easy to create an abstract type that does the scalar-to-broadcasted conversion (a natural/usual direction) automatically:

abstract type BroadcastableCallable 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::BroadcastableCallable) =
    Broadcast.broadcasted(
        constructor_of(typeof(obj)),
        map(Broadcast.broadcastable, fieldvalues(obj))...)

call(f, args...) = f(args...)

Broadcast.broadcasted(c::BroadcastableCallable, args...) =
    Broadcast.broadcasted(call, c, args...)

(You can think of BroadcastableCallable as a type that does struct-of-arrays -to- array-of-structs conversion lazily/on-the-fly.)

You can then just create callable subtypes

struct Scale{T} <: BroadcastableCallable
    a::T
end

struct Shift{T} <: BroadcastableCallable
    a::T
end

struct Compose{F, G} <: BroadcastableCallable
    f::F
    g::G
end

(f::Scale)(x) = f.a * x
(f::Shift)(x) = f.a + x
(f::Compose)(x) = f.f(f.g(x))

which has the desired properties I described; i.e., it works with scalar parameters:

julia> f = Compose(Scale(10), Shift(0.1));

julia> f.([1, 2, 3])
3-element Array{Float64,1}:
 11.0
 21.0
 31.0

as well as vector parameters:

julia> f = Compose(Scale([10, 100, 1000]), Shift([10, 20, 30]));

julia> f.([1, 2, 3])
3-element Array{Int64,1}:
   110
  2200
 33000
1 Like