Lazy broadcasting macro

[mcabbott]

This PR adds a macro @lazy to turn broadcasting expressions into BroadcastArrays, and @lazydot which applies @. first. (It also replaces #18 with an explicit sum method.)

The clever idea here to hack broadcasting is due to @dawbarton and @tkf, discussed here:

https://discourse.julialang.org/t/boolean-short-circuit-fusion/19641/20
JuliaLang/julia#19198 (comment)
JuliaLang/julia#30939

Perhaps this package would prefer an @lazy which also converted vcat -> Vcat etc, haven't thought about that. The name @lazy also clashes with MacroTools Lazy.jl; I was going to suggest @z & @ż as short names but can't seem to type ż at the REPL (on a mac)... how ugly are @□ & @⊡?

[dlfivefifty]

Great! Though I think we should think through the names of the macros carefully. Once Applied is merged, we have a few different "lazy" objects and so @lazy on its own may be confusing:

1. broadcasted in base. I don't think we need to worry about this one.
2. BroadcastArray
3. applied in LazyArrays.jl. Plays the same role as broadcasted but for lazy applying. E.g., applied(exp, A) is an un-realised matrix exponential.
4. ApplyArray. For example, ApplyArray(exp, A) represents the matrix exponential as an AbstractMatrix.

[mcabbott]

Yes, this is just a first attempt, I was playing with macros and found myself in danger of re-inventing the wheel... better to use BroadcastArray.

Agree the name is confusing, or too limited, in the context of the rest of this package. One could conceivably make @lazy exp(A) == ApplyArray(exp, A) and @lazydot exp(A) == applied(exp, A) but I'm not sure I volunteer to write such a thing. The more limited actually existing one could be called @lazycast? @bc? I don't know.

Not exported for now, but this is implemented as lazy.(exp.(A)) which is a little odd, but at least clear it has something to do with broadcasting. It has also been suggested that this could be a unary operator, instead.

[codecov-io]

Codecov Report

Merging #21 into master will decrease coverage by 0.49%.
The diff coverage is 36%.

@@            Coverage Diff            @@
##           master      #21     +/-   ##
=========================================
- Coverage   50.06%   49.56%   -0.5%     
=========================================
  Files          11       11             
  Lines         785      809     +24     
=========================================
+ Hits          393      401      +8     
- Misses        392      408     +16

Impacted Files	Coverage Δ
src/lazybroadcasting.jl	39.58% <36%> (-6.25%)	⬇️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 8f2e313...c507342. Read the comment docs.

[tkf]

How about using a supre short macro like @~? It would be less distracting in an equation and we can "avoid" the naming problem. Tilde kind of look like a wave and so I think it can mean "floating", i.e., non-materializing behavior.

Once Applied is merged, we have a few different "lazy" objects

Any reason why for exposing both the expression trees (broadcasted and applied) and array wrappers (BroadcastArray and ApplyArray) to the end-users (i.e., not library writers extending LazyArray)? If we don't immediately materialize the expression (e.g., Y .= Mul(A, X)) I think it's nicer to have an array-like object so that it can work nicely with the rest of Julia ecosystem. As for BroadcastArray vs ApplyArray, I don't think we should choose one. It's better if I can write @lazy (A .* B) ⋆ C.

[tkf]

Isn't esc(:( $lazy.($ex) )) better? It'd work better with macros inside @lazy.

[mcabbott]

I never thought to use $ for that, am new to macros. Both give @macroexpand @~ @. A + 2 --> :((LazyArrays.lazy).((+).(A, 2))), how do I see the difference?

[tkf]

Ah, sorry, never mind. I think I was mixing it with something else.

[tkf]

Another point: why do you need @lazydot? Why not @lazy @.?

[mcabbott]

Yes it could be shorter, I wondered about having a pair of symbols with & without a dot, but don't much like the ones I could think of above. I like the @~ idea, I never write ~ so this suggests nothing else to me. Indeed there is no great need for a separate @lazydot.

[dlfivefifty]

I like it! One could imagine ~ playing the equivalent role as ., eg y ~= exp~(A) could in the future be in-place matrix exponential.

[tkf]

I should have mentioned that the downside of ~ is that it already means bitwise not. If you are somehow using ~x all the time, it may be confusing. But it seems you guys like it so maybe I don't need to worry :)

[mcabbott]

Well spotted, Base.:~(f::Function, x::Any) = Applied(f,x) seems to work.

[dlfivefifty]

Let’s avoid inconsistency if we can. I think a Unicode symbol is better than overloading bitwise not.

[tkf]

An ascii symbol that is available and not an operator is $. So, we can use @$ in principle. There is also @_. It would be nice if we can use @: but that's a parse error...

[mcabbott]

I've never written ~ intentionally; someone who does would have to tell me if it ever happens around broadcasting. Other short names: @_ (also downstairs?) @¬ (also a negation, \neg) @| (barrier to materialisation) @□ (open \square)?

Anyway no rush, this is mostly to see if it functions well first. My tests like #16 are about large arrays, but perhaps other things (such as allocations of small broadcasts) may need more scrutiny.

[tkf]

@.. is also usable :)

[dlfivefifty]

Ideally it would look "nice" as dot syntax. Not that it's going to be incorporated into the parser anytime soon, but just better to have it evoke it. .. is OK:

y ..= exp..(A)

but I think it's more appropriate for a "tiered" broadcast (broadcast(v -> broadcast(exp, v), A)).

□ is OK:

y □= exp□(A)

but not great.

¬ is also OK:

y ¬= exp¬(A)

but knowing AppleScript I just see "new line".

_, |, and : are not good.

[tkf]

but I think it's more appropriate for a "tiered" broadcast (broadcast(v -> broadcast(exp, v), A)).

I (somewhat) agree about this point. But I note that f..(x) is actually previously considered (JuliaLang/julia#16285 (comment)) but I think it's now discarded as it is already a valid syntax (JuliaLang/julia#19198 (comment); you can see that it's :(f..(x)) == Expr(:call, :.., :f, :x)). So, I don't think it's super crazy to use @.. here as it's already considered before separately.

[dlfivefifty]

Ah, we shouldn't use .. as a .. b is the interval a ≤ x ≤ b in IntervalSets.jl!

[tkf]

Any reason why for exposing both the expression trees (broadcasted and applied) and array wrappers (BroadcastArray and ApplyArray) to the end-users (i.e., not library writers extending LazyArray)?

Answering to myself: BroadcastArray and ApplyArray need inference. In the middle of high-performance code, it's better to avoid it if possible. So, I think returning broadcasted and applied makes more sense. We can have a factory function (say) array to wrap them into BroadcastArray and ApplyArray; e.g., array(@lazy exp((A .* B) ⋆ C)).

[dlfivefifty]

Why not just convert(AbstractArray, @lazy ...) ?

[tkf]

That makes sense too. I also actually thought about AbstractArray(@lazy ...). But I liked a shorter name.

[dlfivefifty]

I think we should keep things minimal for now and avoid adding names if possible. I think AbstractArray(...) should also work.

[tkf]

OK. Avoid introducing new names makes sense.

[mcabbott]

I'm quite keen on writing sum(@~ A .+ B .* B'; dims=2), this was my reason for wanting a macro which made a BroadcastArray, not just a Broadcasted as in the earlier discussion (19198 etc).

Having to write sum(AbstractArray(@lazy A .+ B .* B')), dims=2) seems clumsy, would this gain anything in terms of inference? I don't understand much but AbstractArray{T,N} implies eltype must be known I guess. Could the pure Broadcasted object be made efficiently sum-able without doing this?

I presume it's not just me, and that not creating a big array you are about to reduce (over some dimensions) is a big reason you'd want broadcasting.

[tkf]

Could the pure Broadcasted object be made efficiently sum-able without doing this?

I believe so. Actually, using broadcasted and applied is likely be more efficient for "not so large" arrays as the compiler can eliminate allocations for those object. I think allocation of the BroadcastArray object cannot be eliminated since it invokes the inference API (and hence not type stable).

Efficient sum would be implemented via mapreduce. This means that the we need to send a PR to Julia. But I think this is good since we get better performance in all Julia ecosystem even outside LazyArray.jl.

[mcabbott]

OK, so something like adding a method mapreducedim!(f, op, R::AbstractArray, A::Broadcasted), and all the steps between sum and that? But one of those needs to create the output array. (Although sum!(C, @~ A .+ B .* B') would also be useful.)

For the complete reduction, chethega's airsum split off s = x[1] first. Does that count as avoiding inference? This stuff may be above my pay grade...

[tkf]

By "need inference", what I meant was "need to call inference API", i.e., this:

LazyArrays.jl/src/lazybroadcasting.jl

Line 17 in 8f2e313

_BroadcastArray(bc::Broadcasted) = BroadcastArray{combine_eltypes(bc.f, bc.args)}(bc)

It looks like combine_eltypes itself can actually be inferred (e.g., try @code_warntype Broadcast.combine_eltypes(+, ([0.0], [1], [missing, 1im]))) although I don't know if it fails in more complex ones. So I probably don't need to worry too much. But using inference API is something advised to be avoided so I'd say it's safer to not use it. Another argument is that you can opt-in array version by wrapping the lazy macro with AbstractArray but the other way around is impossible.

[tkf]

I found a horrible hack to define a macro @:: 😆

julia> ex = :(macro _(ex); :(lazy.($(esc(ex)))); end);

julia> ex.args[1].args[1] = :(::)
:(::)

julia> eval(ex)
@:: (macro with 1 method)

julia> @macroexpand @:: a .+ b
:(Main.lazy.(a .+ b))

[nalimilan]

It looks like combine_eltypes itself can actually be inferred (e.g., try @code_warntype Broadcast.combine_eltypes(+, ([0.0], [1], [missing, 1im]))) although I don't know if it fails in more complex ones. So I probably don't need to worry too much. But using inference API is something advised to be avoided so I'd say it's safer to not use it. Another argument is that you can opt-in array version by wrapping the lazy macro with AbstractArray but the other way around is impossible.

See also JuliaLang/julia#30485.

But indeed it's not recommended to expose information obtained through inference to the user (e.g. via the element type of the array). Since the element type of broadcast depends on the actual contents of the input data (e.g. whether there are missing values or not), it cannot be type stable. Inference can only compute an upper bound unless the input type is concrete (neither abstract nor Union).

What happens in practice if you don't define the eltype of BroadcastArray, or always return Any? Maybe that would work?

[dlfivefifty]

OK, we have the following different cases:

1. Materialized broadcasted. With macro @.
2. Lazy broadcasted. Not exposed in Base.
3. BroadcastArray. No macro has been proposed, but could be @a.
4. Materialized applied.
5. Lazy applied. Macro proposed @~.
6. ApplyArray. Macro could be @a~.

So there's a case for macros for 3,5, and 6. What about 2 and 4?

[mcabbott]

I like the list. The present macro @~ of this PR is 3. right? And 2. is the @lazy from the discourse thread. Does 4. mean that Broadcast.materialize is lazily applied somehow?

Just to have toy examples somewhere, this is the behaviour with some unions. I don't know how objectionable these are, BroadcastArray{Union{... seems preferable to BroadcastArray{Any,... surely?

julia> ints = [2,3,4];
julia> maybe = Union{Int,Missing}[5,6,7];
julia> f_un(x) = rand() < 0.01 ? missing : x;

julia> ints .+ maybe' # 3×3 Array{Int64,2}

julia> mat = @~ ints .+ maybe' # 3×3 BroadcastArray{Union{Missing, Int64}, ...

julia> collect(@~ int .+ maybe') # 3×3 Array{Union{Missing, Int64},2}; same for copy(), Array()

julia> sum(int .+ maybe'; dims=1) # 1×3 Array{Int64,2}

julia> sum(@~ int .+ maybe'; dims=1) # 1×3 Array{Union{Missing, Int64},2}

julia> f_un.(ints) # 3-element Array{Int64,1}, usually! 

julia> @~ f_un.(ints) # 3-element BroadcastArray{Union{Missing, Int64}

I think we should make collect call Broadcast.materialize here, there should be an easy way to do that. It's possible that copy should too, e.g. for lazy transpose it makes a full Array.

Setting maybe[2] = missing will change mat thus it can't have a narrower type.
But would it be possible to make sum(::BroadcastArray, dims=1) magically (unstably) produce Array{Int64}? That ones seems sub-optimal.

[mcabbott]

ps. That hack for @:: also lets you define @..., so 1 to 4 on the list are covered!

julia> ex.args[1].args[1] = :(...)
:...

julia> @macroexpand @... a .+ b
:((Main.lazy).(a .+ b))

[dlfivefifty]

Does 4. mean that Broadcast.materialize is lazily applied somehow

In the PR #17, for DefaultApplyStyles, a materialized applied is the same as just calling the functions:

materialize(applied(f, applied(g, x, y)) === f(g(x,y))

But there can be other ApplyStyles that materialize in a different way. This is used for Mul (which corresponds to applied(*, ...)) so that the materialize is based on MemoryLayout, instead of type, thereby avoiding the issues with mul!. I.e.

materialize(Mul(A,B)) === materialize(applied(*, A, B)) ≠ A*B

Another example where this is useful would be matrix-exponentials times vectors. That is,

materialize(applied(*, applied(exp, A), v))

could use the fact that exp(A)*v can be done without constructing exp(A).

[tkf]

What happens in practice if you don't define the eltype of BroadcastArray, or always return Any?

@nalimilan I think the performance would be bad. For example, similar would return an Array{Any} which would affect any functions that has to allocate buffer or output collection. I think having a basic set of API not relying on inference (broadcasted and applied) and a high-level easy-to-use API based on inference (BroadcastArray and ApplyArray) is a good balance.

[tkf]

OK, we have the following different cases:

@dlfivefifty I think it's better to have just one macro that composes well. That is to say, I suggest to implement @~ that recursively replaces all dot calls with broadcasted and all non-dot calls with applied. By recursively I mean that apply the same rule to dot calls inside non-dot calls and vice versa. That is to say, @~ (A .* B) * C .+ D would be lowered to

x = broadcasted(*, A, B)
y = applied(*, x, C)
broadcasted(+, y, D)

Note that this recursive nature is compatible with how @. applies the dot to all calls.

With this approach, we cover all the list by composing @~ and @.:

1. Materialized broadcasted: @. EXPR
2. Lazy broadcasted: @~ @. EXPR
3. BroadcastArray: AbstractArray(@~ @. EXPR)
4. Materialized applied: lhs .= @~ EXPR or copy(@~ EXPR)
5. Lazy applied: @~ EXPR
6. ApplyArray: AbstractArray(@~ EXPR).

Note: if you want to reduce the length by 1, you can also write @~@.

I know that 3 and 6 are rather long. However:

• You'd want to have an array when you are at "API boundary"; i.e., you are calling functions that are not aware of broadcasted or applied. In this case, it's better to be explicit if you want a lazy or materialized array. For example, you better materialize the array before passing it to *.

• We don't need a BroadcastArray for invoking mapreduce/mapfoldl if we "fix" them in Base. I think this covers most of the cases for lazy reduction (e.g., sum).

• Importantly, 3 and 6 do not require macro.

I think we should avoid using macros whenever possible. However, we better define a really good one if we need to do so. The recursive nature of what I suggest for @~ is harder to implement. But I think it is worth it since it's simple to use.

[tkf]

FYI I started implementing mapreduce for Broadcasted. sum(Broadcast.broadcasted(*, a, b)) is now 5x faster than before and almost as fast as dot. JuliaLang/julia#31020

[mcabbott]

This plan sounds great, if just a little more ambitious than my PR! Is the idea that @~ would be part of this package, or would be extended somehow to add 3-6?

[tkf]

Yes, my idea for @~ makes sense only if @~ lives in LazyArrays.jl since it needs applied which is defined in LazyArrays.jl.

[dlfivefifty]

There is a case for applied to live in its own lightweight package (LazyOperations.jl?) separate from lazy arrays, but I think that's not terribly important right now.

I'll try to finish the Applied PR the next few days.

[nalimilan]

@nalimilan I think the performance would be bad. For example, similar would return an Array{Any} which would affect any functions that has to allocate buffer or output collection. I think having a basic set of API not relying on inference (broadcasted and applied) and a high-level easy-to-use API based on inference (BroadcastArray and ApplyArray) is a good balance.

Right. But maybe that just means it shouldn't support similar, as it's not possible to determine the type of the result without relying on inference (which really isn't recommended).

[dlfivefifty]

I’m not sure I understand, Broadcasted supports similar without issue.

[nalimilan]

Not AFAICT:

julia> similar(Base.Broadcast.broadcasted(+, [1]))
ERROR: MethodError: no method matching similar(::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1},Nothing,typeof(+),Tuple{Array{Int64,1}}})

You need to provide the element type:

julia> similar(Base.Broadcast.broadcasted(+, [1]), Int)
1-element Array{Int64,1}:
 0

[dlfivefifty]

Right, but one could easily define a similar(::Broadcasted):

julia> bc = Broadcast.broadcasted(+, [1]);

julia> similar(bc, Broadcast.combine_eltypes(bc.f, bc.args))
1-element Array{Int64,1}:
 0

[nalimilan]

combine_eltypes depends on inference, so it can return Any if inference fails:

julia> bc = Broadcast.broadcasted(x -> x > 0 ? (x > 1 ? (x > 2 ? nothing : 1) : 1.0) : "a", [1]);

julia> Broadcast.combine_eltypes(bc.f, bc.args)
Any

[tkf]

3. BroadcastArray: AbstractArray(@~ @. EXPR)

Talking to myself: This actually is a type piracy. So maybe it's better to have a factory function array so that array(@~ @. EXPR) works.

[mcabbott]

3. BroadcastArray: AbstractArray(@~ @. EXPR)

How does this differ from BroadcastArray(@~ @. EXPR)?

2. @~ @. EXPR

If sum(::Broadcast; dims) ends up in Base, then perhaps the this level 2 macro should join it? With the stronger macro of 4,5,6 living here (tell people that @≈ is two waves?)

I also wonder if there is room for confusion between these levels. If @~ f(A) is level 5, lazy applied, then does @~ f(A) .+ B make a lazy application of f in addition to lazy broadcasting?

[tkf]

How does this differ from BroadcastArray(@~ @. EXPR)?

It's nice to be able to write array(x) where x can be a broadcasted or applied.

If sum(::Broadcast; dims) ends up in Base, then perhaps the this level 2 macro should join it?

I think it would go in to Base as soon as the discussion in JuliaLang/julia#19198 is settled. I'm trying to speed up it a little bit by making mapreduce faster in JuliaLang/julia#31020.

I also wonder if there is room for confusion between these levels. If @~ f(A) is level 5, lazy applied, then does @~ f(A) .+ B make a lazy application of f in addition to lazy broadcasting?

Since it is syntactically clear that f(X) is a call and f(A) .+ B is a dot-call (with a call inside its argument), I don't think there is any confusion. Also, in the end, you just get some thing that can be materialized via lhs .= thing for copy(thing). I think it's nice to treat them uniformly.

[dlfivefifty]

I really don't like array as its too generic and could be confusing. We could do:

abstract type LazyArray{T,N} <: AbstractArray{T,N} end
struct BroadcastArray{T,N} <: LazyArray{T,N} end
struct ApplyArray{T,N} <: LazyArray{T,N} end

then support LazyArray(@~ @. EXPR)

[tkf]

I really don't like array as its too generic and could be confusing.

This makes sense. LazyArray(@~ @. EXPR) sounds good to me.

[tkf]

Here is a sample implementation of what I proposed in #21 (comment) which can do something like

julia> @macroexpand @~ a .+ f(b, c)
:((Base.Broadcast.instantiate)((lazy).(a .+ (LazyArrays.applied)(f, b, c))))

julia> @macroexpand @~ a .+ f(b, c .+ d)
:((Base.Broadcast.instantiate)((lazy).(a .+ (LazyArrays.applied)(f, b, (Base.Broadcast.instantiate)((lazy).(c .+ d))))))

Code:

using LazyArrays

lazy(::Any) = throw(ArgumentError("function `lazy` exists only for its effect on broadcasting, see the macro @~"))
struct LazyCast{T}
    value::T
end
Broadcast.broadcasted(::typeof(lazy), x) = LazyCast(x)
Broadcast.materialize(x::LazyCast) = x.value


is_call(ex::Expr) =
    ex.head == :call && !startswith(String(ex.args[1]), ".")

is_dotcall(ex::Expr) =
    ex.head == :. || (ex.head == :call && startswith(String(ex.args[1]), "."))
# e.g., `f.(x, y, z)` or `x .+ y .+ z`

lazy_expr(x) = x
function lazy_expr(ex::Expr)
    if is_dotcall(ex)
        return bc_expr(ex)
    elseif is_call(ex)
        return app_expr(ex)
    else
        # TODO: Maybe better to support `a ? b : c` etc.? But how?
        return ex
    end
end

function bc_expr(ex::Expr)
    @assert is_dotcall(ex)
    return :($(Broadcast.instantiate)($lazy.($(bc_expr_impl(ex)))))
end

bc_expr_impl(x) = x
function bc_expr_impl(ex::Expr)
    # walk down chain of dot calls
    if is_dotcall(ex)
        return Expr(ex.head,
                    lazy_expr(ex.args[1]), # function name (`f`, `.+`, etc.)
                    bc_expr_impl.(ex.args[2:end])...) # arguments
    else
        return lazy_expr(ex)
    end
end

function app_expr(ex::Expr)
    @assert is_call(ex)
    # instantiate?
    return app_expr_impl(ex)
end

app_expr_impl(x) = x
function app_expr_impl(ex::Expr)
    # walk down chain of calls and lazy-ify them
    if is_call(ex)
        return :($applied($(app_expr_impl.(ex.args)...)))
    else
        return lazy_expr(ex)
    end
end

macro ~(ex)
    # checkex(ex)
    esc(lazy_expr(ex))
end

[dlfivefifty]

Let me know when this is ready to merge.