Overloading vectorization?

So say I have a function, e.g. f(x,y) = x*y and two arrays a=[1,2]; b=[3,4], then I can perform f element-wise by calling f.(a,b).

But now f is a more complicated operation, f.(a,b) still works, but say I can write another version of the function g such that
g(a,b) == f.(a,b)
but g is much more efficient than the element-wise operation. Can I overload f. to perform g instead of an element-wise f? So a user, not knowing the difference, doesn’t get the bad performance by naively using the f. version.

1 Like

Can you give an example of the type of f you are worried won’t vectorize efficiently?

1 Like

Vectorisation with . is just syntactic sugar for broadcast(f, args..)

So you can absolutely override that. But the broadcast mechanism is complicated. You would need to dispatch on a specific subset of broadcast for the function f operating on some specific types. Add methods to Base.broadcasted, as @mbauman says below. I’ve never tried that but I guess it’s possible.

There might be other ways to do this too, but you’ll have to learn a bit about broadcast.

2 Likes

Sure, I have a piecewise linear function on a grid of nodes, i.e I know the grid points and the values. Then I want to evaluate this at some arbitrary point in the confines of the grid. To do this I go over the nodes to find the 4 grid points representing the corners of the square and compute the interpolated value as a weighted sum of the nodal values.

Now I have instead a list of say millions of points I want to evaluate in the confines of the grid. I could do the process above for each, but that would make me traverse the grid nodes for each point. It is many many times more efficient to compute the nodes simultaneously for all points which can be done in a single traversal.

Okay, thanks, I will have to take a look at that.

That’s not true. It’s a sugar for a much more complicated piece of infrastructure:

julia> Meta.@lower f.(x, y)
:($(Expr(:thunk, CodeInfo(
    @ none within `top-level scope'
1 ─ %1 = Base.broadcasted(f, x, y)
│   %2 = Base.materialize(%1)
└──      return %2
))))

You can overload Base.broadcasted(::typeof(f), x::Vector, y::Vector).

This is the mechanism through which, for example, ranges use O(1) arithmetic and recompute new ranges when possible.

2 Likes

Errr… it’s true in so far as the docs say pretty much what I said:

A special syntax exists for broadcasting: f.(args…) is equivalent to broadcast(f, args…), and nested f.(g.(args…)) calls are fused into a single
broadcast loop.

Which is a useful start for someone who doesn’t know that.

I also mention that broadcast is complicated, and that there probably is a better way to do it.

Yeah, there’s a tension between documenting the behavior for calling users and the behavior as you overload it… the documentation there is correct for how broadcast works when you call it but it doesn’t translate to how dot-call extension works.

It may seem like splitting hairs, but what you wrote was very different. I’m sorry my message was so short and curt — just wanted to point the way to broadcasted rather than broadcast.

2 Likes

No problem! It was the wrong method lol. I was just trying to point the way to broadcast in general, as it hadn’t been mentioned… but I can honestly never remember how it works even after implementing it in a well used package :grimacing:

As a note for anyone reading this later…

If f is a struct of type S so that f(x,y) is defined as (s::S)(x,y) = ... instead of “::typeof(f)” in Base.broadcasted(::typeof(f), ...) you simply use s::S so you can refer back to the object itself as s in evaluation.

1 Like