Intelligent broadcasting

[lamorton]

Branching from #128, I'd be interested to see AxisArrays provide a method to automatically align AxisArrays for broadcasting in a manner parallel to what xarray does in Python.

For those of you who aren't familiar, the final line of this:

A = AxisArray(randn(10,15), :x, :t)
B = AxisArray(randn(10,20), :x, :y)
C = AxisArray(reshape(A, (10,15,1)) .* reshape(B, (10,1,20)), :x, :t, :y )

would become this:

C = A .* B

which is especially nice when one wants to abstract over the number shared & non-shared dimensions. Presently, the cleaner syntax here would fail due to the mismatch between 15 & 20. Additionally, even

C = A .* A

presently produces C as a plain array, not an AxisArray, which discards the axis information.

There are some gotchas: the ordering of the resulting dimensions could matter for performance reasons, and I don't think xarray has a way to specify how this is done. Additionally, if there are vectors associated to the axis names, then one has to decide if/how to test equality for same-name axes on different arrays, and what to do in the case of inequality. (xarray chose pointwise equality tests and silently drops any non-matching points along each axis. While I understand the motivation to imitate a database join, I've been bitten by floating-point comparisons in this context.)

[junglegobs]

I would really appreciate this too! I don't care too much about performance in this case, since I use my AxisArrays as (much nicer) alternatives to Dict{Tuple,Number}. EDIT: Surely this is fairly simple to implement? It's a fairly simple redefinition of Base.broadcast right?

[iamed2]

EDIT: Surely this is fairly simple to implement? It's a fairly simple redefinition of Base.broadcast right?

No, it's actually quite challenging. You can't effectively overload broadcast for some set of types, you need to interact with the internals of broadcast, and there are many levels. It's not implemented there yet, but perhaps @oxinabox has thought about how to do this for NamedDims.jl?

[oxinabox]

I think I have a way, scroll to the bottom. This response is stream of consciousness. That is what happens when I decide to answer questions at 1am.

---

I have thought about automatic alignment a lot. Largely about if we do it at all, which is not clear to me. (and secondly about holw to handle the wildcard names, whch s not applicable here)

It would be notably more complex than NamedDims current broadcasting. Which is fairly complex https://github.com/invenia/NamedDims.jl/blob/51893cfd36a771061c9fc0cf490d04ad5c7d86a1/src/broadcasting.jl (Though it is actually almost as simple as iit gets for wrapper Array types -- they basically all need to do this or they can't suport static arrays).

Right now NamedDims gets to:

• strip all the names off at the start,
• rewriite the Broadcasted expressions to only have backing arrays and not to use the special style
• to the broadcast on the simplified expression
• Workout what the names of the results get (and if it is legal).

I don't know that tranposing is legal in broadcasting. Thing is that broadcasting never changes the shape of the data. All the datas (arrays/scalars) in a fused broadcasted operations have the same shape. If we treat them as having shape right padded with infinte 1s. I.e. scalars have shape (1, 1, 1, ....), vectors have shape (length(xs), 1, 1, 1,...) arrays in general have (size(xs).., 1, 1, 1...)

So every time you did an automatic tranpsosition, you would need to break the broadcast fusion at least. Which I guess is possible -- I mean TensorFlow.jl breaks broadcast fusion down to not fuse for every single binary operator. I guess it should be possible to only split the fusion if a transposition would occur.

Hmm, maybe better:

• Precompute all the transpositions that well be needed at the start on the whole Broadcasted expression. (that means in copy and in copyto!
  • Rewrite the Broadcasted expression to be one that: A) doesn't use the special style, and B) has already had the transpositions performed
  • Do the broadcast on the simplified expression
  • put all the names an indicies back on at the end.

I think this is valid, not really all that different to what NamedDims does already.

[c42f]

they basically all need to do this or they can't suport static arrays

That's a worry. Is this because of something nasty we're doing in StaticArrays, or due to a limitation of the Base API (or both)? I think the StaticArrays support for broadcast is somewhat inherently complex but if it doesn't compose with any other packages that's a bit sad.

[oxinabox]

Static Arrays

1) has its own broadcast style so needs to ensure that styles wrap existing styles that's most of the complexity and many other packages' arrays need it too.

2) Static Arrays are not mutable so just implementing similar like the Base docs say is not enough. Need to do copy instead. (Tracker also needs this)

The broadcast machinery and docs were not designed with composition in mind I think.

[c42f]

Ok thanks a lot, those sound like things which eventually need fixing in Base (to add to the many other places where the immutability of SArray breaks assumptions of the Base array code...)

[mcabbott]

Perhaps of interest, I had a go at making something like this work, and now it's registered:

using NamedPlus
A, B = ones(x=10, t=15), ones(x=10, y=20);

@named C{x,t,y} = A .* B  # "10×15×20 NamedDimsArray(::Array{Float64,3}, (:x, :t, :y))"

What this ends up doing is transmute(A, (1, 2, 0)) .* transmute(B, (1, 0, 2)), where transmute is a more flexible version of PermutedDimsArray, which here just inserts trivial dimensions. (From TransmuteDims.jl.) I don't think there's any need to break dot fusion.

The "permutations" required are worked out from the names at compile-time. At the moment you have to explicitly state what names the output should have, and in what order, although I suppose it would not be hard to make it guess.

It's obviously much easier to make a macro do all of this, than to make the broadcasting machinery do it. Like @oxinabox above, I'm not certain this should be the default, as it seems like it could add a lot of surprise. For instance v .+ v' would presumably then make a vector.

[lamorton]

@mcabbott "Explicit is better than implicit," so +1 for a solution that makes it possible to declare the output dimension tuple.

What happens to 'dimensional axes'? It looks like the AxisArray default behavior is to return a bare array from A .* B.

[mcabbott]

I think AxisArrays never got its broadcasting updated for Julia 1.0, the old machinery was much simpler (but less powerful). It could surely be done though.

My @named macro only works for NamedDimsArrays. It strips off the names completely before broadcasting & puts them back afterwards. So I am sure it (or something like it) could be made to work for AxisArrays. There would be a little more work in that it would have to preserve not just the names but the "axis vector" things.