Support new_axes= keyword in atop

[mrocklin]

Add new single-chunk dimensions with the new_axes= keyword, including
the length of the new dimension. New dimensions will always be in a
single chunk.

>>> def f(x):
...     return x[:, None] * np.ones((1, 5))

>>> z = atop(f, 'az', x, 'a', new_axes={'z': 5})

[mrocklin]

This could use high-level review on API. Can anyone think of a cleaner way to do this or attractive alternatives?

@jcrist @shoyer

[shoyer]

This looks like a slightly less general version of #1511? The ability to insert new axes is certainly more important than resizing chunks (I can no longer remember exactly what use case I had in mind there).

[mrocklin]

Hrm, indeed. I had forgotten about #1511 .

The solution presented here does less, but the changes are also more modest, which reduces the concerns about maintenance bloat a bit (though not to zero).

I think the relevant question now is "Are we likely to want multi-chunk new dimensions?" If the answer is "yes" then we should think harder about this API.

I'm actually unsure what adding a new multi-chunk dimension would look like. The user defined function doesn't have a clear way to output multiple chunks. Thoughts?

[shoyer]

The solution presented here does less, but the changes are also more modest, which reduces the concerns about maintenance bloat a bit (though not to zero).

You also figure out a more elegant/minimal way to adjust top than I did :).

I think the relevant question now is "Are we likely to want multi-chunk new dimensions?" If the answer is "yes" then we should think harder about this API.

My version didn't actually support mulit-chunking new dimensions. It supported changing the chunking of existing dimensions on the output while keeping "block indices" intact. This is a pretty natural sort of thing to do (adjusting the size of each chunk), but maybe more complex than warranted without clear use cases.

[shoyer]

This (and concatenate) should go after *args, and be marked as "keyword only"

[mrocklin]

Fixed

[mrocklin]

My version didn't actually support mulit-chunking new dimensions. It supported changing the chunking of existing dimensions on the output while keeping "block indices" intact. This is a pretty natural sort of thing to do (adjusting the size of each chunk), but maybe more complex than warranted without clear use cases.

Ah, I see. Is there a way to combine that use case into this one or a way to make this change future proof so that it doesn't conflict with future desired changes?

[shoyer]

I don't think there's any conflict here. Every valid argument to your new_axes was also a valid argument for my out_chunks, but new_axes is certainly a more intuitive name. I think this can go in as is.