implement `@broadcast_dims` macro

[haakon-e]

I've really enjoyed broadcast_dims since discovering it a few months ago.

Using it a lot, I've found that for some operations, the broadcast_dims notation can be a bit obfuscating, and that macros can be used to simplify code that may previously have been relatively complex.

If there's appetite for something like this, I'll open a PR, with a few tests and docs when I have time.

A few examples:

xs, ys, zs = (X(-180:2.0:-150), Y(-30:2.0:-10), Z(-5000:100.0:0))

a = rand(xs, ys, zs)

d = rand(ys, zs)

# currently, I have to do:
broadcast_dims(/, a, d)

# with my proposal, below, I can do:
@broadcast_dims a / d

# it's implemented to be arbitrarily chained, with certain calls left untouched, e.g.
@macroexpand @broadcast_dims a + .√a / (sqrt.(d) * a.^2)
# expands to:
:(broadcast_dims(+, a, broadcast_dims(/, .√(a), broadcast_dims(*, sqrt.(d), a .^ 2))))

implementation sketch

using MacroTools: @capture  # light dependency: https://github.com/FluxML/MacroTools.jl/blob/master/Project.toml
export @broadcast_dims
macro broadcast_dims(expr)
    function process_expr(ex)
        if @capture(ex, f_(a_)) || @capture(ex, a_.b_) || @capture(ex, f_.(args__))
            # don't touch: (1) single arg functions `f(a)`, (2) field access `a.b`, (3) function broadcast `f.(a, b, c)`
            return esc(ex)
        elseif @capture(ex, f_(args__))
            string(f)[1] == '.' && return esc(ex)  # don't broadcast a call like `.√a`
            # recursively check each argument for sub-expressions. 
            args2 = map(process_expr, args)
            return :(broadcast_dims($(esc(f)), $(args2...)))
        else
            # don't touch: anything else
            return esc(ex)
        end
    end

    return process_expr(expr)
end

PS: I'm pretty new to macros, but I hope this has decently good "hygiene"

[haakon-e]

ahhh lol I see you might have actually done this in #777

[rafaqz]

Yes! It's working well, try it out. There are some examples in the tests.

The trick to the macro is getting every argument into a separate variable so we can extract the dimensions, then permuting them to match before they are used in the broadcast.

It also lets you specify the output dims order too so it's deterministic for whatever inputs.