Specializing broadcasting for a scalar function

Let’s say I have a function

function f(a, b)
     temp = do_something_expensive(a)
     do_other_thing(temp, b)
end

I like writing scalar functions because they don’t obscure the logic with additional loops. But in this case, do_something_expensive can bite me performance-wise if I broadcast over a scalar a and a vector b:

# runs do_something_expensive on a_scalar for every b in b_vector
f.(a_scalar, b_vector)

Can one overload broadcasting for the case where a is a scalar and b is a vector in order to only run do_something_expensive once? Or would it make more sense to make another method for scalar and vector arguments and use that instead? I kind of like to just expose one scalar method to the user and do intelligent things in the background if it’s used in broadcasted style.

1 Like

Yes, this is possible by overloading Base.broadcasted specifically for f:

julia> do_something_expensive(x) = (println("don't run me twice!"); x)
do_something_expensive (generic function with 1 method)

julia> do_other_thing(x, y) = (println("this is fine"); (x, y))
do_other_thing (generic function with 1 method)

julia> function f(a, b)
            temp = do_something_expensive(a)
            do_other_thing(temp, b)
       end
f (generic function with 1 method)

julia> function Base.broadcasted(::typeof(f), a::Number, b::AbstractVector)
           temp = do_something_expensive(a)
           return Base.broadcasted(do_other_thing, temp, b)
       end

julia> f.(1, 1:3)
don't run me twice!
this is fine
this is fine
this is fine
3-element Vector{Tuple{Int64, Int64}}:
 (1, 1)
 (1, 2)
 (1, 3)

Whether this is a good idea is another question though. This approach does make it more obfuscated to see what’s actually going on, so perhaps you just want to have the user call do_something_expensive and do_other_thing themselves.

2 Likes

In this case the user is myself but I see your point :wink: For example f can be a projection function project(scene, point) which needs to precompute some matrix depending on scene, but it shouldn’t do that anew for every point I want to project.

One potential answer is to use memoization on project, but that isn’t always a good solution.

That seems like a reasonable use of specializing Base.broadcasted: it sounds like you just want to optimize the computation without changing behavior. There are lots of examples of specializations like this in Base (e.g. specializing map(vcat, array_of_arrays) for efficient array concatenation).

I think it would be nicer to break this up into

project = projection(scene) # create projector object or function 
p = project(point) 

and then

p = project.(points)

I would prefer that way of organizing the calculation. That way you can create several projectors for different scenes and pass them around, for example.

1 Like