Broadcast of SArray with Array gives Array but we'd like SArray

[c42f]

As noted in #740, we currently have the following annoyance that broadcast with static and normal arrays doesn't result in static arrays:

julia> SA[1,2,3] .+ [1,2,3]
3-element Array{Int64,1}:
 2
 4
 6

I think this can be fixed in principle but the big difficulty with broadcasting is that the output SArray Size may depend on the dynamic size of the Array, for example we cannot infer the size of the second dimension in

julia> SA[1,2,3] .+ [1 2 3]
3×3 Array{Int64,2}:
 2  3  4
 3  4  5
 4  5  6

However, there's many circumstances in which we could make this work: basically any time the StaticArray has all dimensions with Size greater than than 1, broadcast semantics require that the Array which is being broadcast together with it match those dimensions (provided the Array as number of dimensions less or equal to the StaticArray).

I'm not sure how to fix this, but one thought was to try stashing the Size trait within the BroadcastStyle and use that to infer the output shape as part of the recursive reduction of BroadCastStyle(a,b).

@mateuszbaran's HybridArrays.jl may also have some code to deal with this kind of thing (https://github.com/mateuszbaran/HybridArrays.jl/blob/master/src/broadcast.jl)

[mateuszbaran]

HybridArrays.jl doesn't handle this, as this problem has always been very easy for me to work around but a proper solution is not simple. There are many special cases.

[c42f]

So I've spent some time trying to do this. Overall I'm finding that it's really tricky to know exactly what to overload. Broadcast is a complicated beast!

I think the right way is to rework StaticArrays broadcast by overloading Broadcast.instantiate to dispatch on the result of Broadcast.combine_axes, which may be

• Purely SOneTo (result should be a static array)
• Base.OneTo (result is an array)
• A mixture of the above (result is a normal array, or hybrid array if that package is loaded)

I'll try to introduce some new kind of trait to control dispatch so that HybridArrays can do a little type privateering and reuse the StaticArrays broadcast infrastructure.

Some relevant links:

• https://docs.julialang.org/en/latest/manual/interfaces/#man-interfaces-broadcasting-1
• Make broadcast_axes inferrable even if axes are of different types JuliaLang/julia#30074
• Fix inference on axes for broadcasting #547
• fix #32442, broadcasting over non-offset arrays with mismatched axis eltypes JuliaLang/julia#34230

[tkf]

one thought was to try stashing the Size trait within the BroadcastStyle

I thought this would work but I just realized that the problem might be that BroadcastStyle becomes non-associative; i.e.,

BroadcastStyle(
    ::AbstractArrayStyle{2},
    BroadcastStyle(
        StaticArrayStyle(Size(1, 3)),  # StaticArrayStyle embeds Size
        StaticArrayStyle(Size(4, 1)),
    ), # => StaticArrayStyle(Size(4, 3))
)  # => StaticArrayStyle(Size(4, 3))

and

BroadcastStyle(
    BroadcastStyle(
        ::AbstractArrayStyle{2},
        StaticArrayStyle(Size(1, 3)),
    ),  # => ::AbstractArrayStyle{2}
    StaticArrayStyle(Size(4, 1)),
)  # => ::AbstractArrayStyle{2}

yield different results if you want to handle everything in BroadcastStyle.

Is this what happened?

---

Anyway, from this thinking, I think I agree combine_axes seems to be a better place to do this.

[c42f]

BroadcastStyle becomes non-associative

Interesting point, but this isn't what happened! I had an alternative formulation where I converted all AbstractArrayStyle{N} into StaticArrayStyle{Size(Dynamic(), Dynamic(), ...)} when mixed with StaticArrayStyle. I thought the result of that can be consistent though I may have been mistaken (any statically sized dimension with length > 1 overrides Dynamic()).

After doing all that, I realized that Base assigns Style{Tuple} to tuples, so we don't even have the length in the style in that case. This is what really killed the idea because we'd have to rely on combine_axes anyway to handle tuples.

The other point is that it's long been a goal to support more general static axes than SOneTo, so trying to dispatch on axes seems right from that point of view.

[c42f]

By the way, I think the woes with Style{Tuple} are just another good example of the pitfalls of trying to do "too much" computation in the type domain, and using axes will just work out better for reasons I haven't worked out yet ;-) One concrete improvement is that Dynamic() is kind of annoying to use as a placeholder for something which will later be an integer, and I found myself reimplementing all the logic which already exists in combine_axes to get the inferred Size.

[tkf]

I had an alternative formulation where I converted all AbstractArrayStyle{N} into StaticArrayStyle{Size(Dynamic(), Dynamic(), ...)} when mixed with StaticArrayStyle.

I guess this might have worked (when ignoring tuples). I was thinking that you'd want to keep other array styles if you ended up not using the static style. For example, producing a sparse matrix from SA[1] .* spzeros(100, 100).

But I agree about the point that coding with interface where you can "touch" the original object is much easier.

[andyferris]

I’ve definitely moved towards instance traits computed on instances, for example in Dictionaries.jl.

[c42f]

For example, producing a sparse matrix from SA[1] .* spzeros(100, 100).

Yes thanks for bringing this up, I'm not really sure what to do with other array styles in this circumstance. I guess we could wrap them to extract later but that seems a bit heavy handed.

A related issue here is in identifying the output static array type. Currently it's a leftmost rule but that seems fairly arbitrary. Perhaps a promotion-like system would make most sense.

[mateuszbaran]

For example, producing a sparse matrix from SA[1] .* spzeros(100, 100).

Yes thanks for bringing this up, I'm not really sure what to do with other array styles in this circumstance. I guess we could wrap them to extract later but that seems a bit heavy handed.

Maybe mixing SparseArrays.HigherOrderFns.SparseMatStyle with StaticArrayStyle should give SparseMatStyle? We most likely don't want to capture all AbstractArrayStyles here, most likely just DefaultArrayStyle and maybe some other defined by a trait?

[c42f]

Yeah that sounds pretty safe. I think the current version only captures DefaultArrayStyle.

[mateuszbaran]

Yes, currently only DefaultArrayStyle is captured. SparseMatStyle is slightly problematic here since StaticArrays.jl doesn't depend on SparseArrays.

[tkf]

I meant to bring up sparse matrix just as an example of other custom styles. I don't think the solution (whatever it is) should refer to a concrete style (like SparseMatStyle). I think one solution may be to use a two-phase strategy in instantiate

1. Try combine_axes. Return a Broadcasted{<:StaticArrayStyle} if we get static SOneTo for all dimensions.
2. Otherwise, re-run combine_styles recursively pretending all the static arrays are regular arrays.

[c42f]

Hmm, a funny/annoying thing here is that — due to lazy fusing — the inferrability of static axes is not nested in the same way as the syntax tree. That's in contrast to the reduction of broadcast styles which follows the AST.

I think this leads to the desire to rerun combine_styles which feels kinda wrong but practical.

Side note: there's also the option of making all purely-StaticArray broadcasts eager rather than lazy by overriding broadcasted(). I'm not sure that's a good idea though and it doesn't really help the other problems here.

[tkf]

Making broadcasted eager sounds like a good approach in terms of inferrability. But I wonder if it makes compiling a .* b .+ c into fused multiply-add harder.

Another disadvantage may be that it makes impossible to create computation tree using, e.g., LazyArrays.@~ when you have static arrays. I don't have a concrete usecase in mind, though.