Add a newsplit argument to transpose

[freeman-lab]

Currently, transpose is significantly less general than swap because it keeps the the split the same. For example, if we want to go from a (1000) x (500,500) array to a (500,500) x (1000) array, we can do that with a swap via swap((0),(0,1)), but with transpose we can only get to (500) x (500,1000) via transpose(2,1,0).

If transpose supported a newsplit or similar keyword arg, we could get to the same result with transpose, which is often more intuitive than calling swap directly.

cc @jwittenbach

[shoyer]

My two cents on API design is that transpose should not move any actual data around, e.g., transpose on (1000) x (500, 600) should yield (600, 500) x (1000), with the key moved from the first to the last axis. Swap operations are expensive, and generally should only be triggered explicitly, or if you're absolutely sure they are necessary for efficient computation. This is particularly relevant for operations like transpose, which based on my experience with NumPy (and dask.array) I expect to be basically free.

[freeman-lab]

Thanks for the feedback @shoyer !

Transpose could behave as you described, it would mean changing the default as to where the split goes -- currently the default is to keep the split the same, you're suggesting a new split that keeps the same contents on either side. But your version would also in general require a swap via shuffle; there's no placement of the split that will make that transpose free, it'll only be free if transposing happens locally to the keys, or the values, or both.

Great insight that people might expect transpose to be more or less free.

One option would be to require manually specifying that swap=True during a transpose, and throw an error if the desired result can't be achieved without a swap.

Also, note that transpose is lazy even if it incurs a swap, which means it will get combined with downstream distributed operations and executed all at once.

[shoyer]

Maybe the bolt array data model should use an explicit list of axes in the key, rather than a split? That would let you guarantee that operations like transpose never need to move around data.

Xray will actually transpose automatically to align named axes in binary operations, so we do this sort of thing a lot.

[shoyer]

Another reason to consider splitting along an explicit set of axes would be because the notion of axis order is not so well defined for some xray objects (datasets can have arrays with different dimensions and axis orders).

I do understand how a split makes sense in the context of broadcasting operations, because the last axes (where broadcasting starts) are not the distributed ones.