Consider adding alternative axis argument to `unpack`

[ricardoV94]

Description

unpack is "awkward" in that it design assumes it should come after pack, and tries to work so that the user can pass the same keep_axes argument:

p, shapes = pt.pack(x, y, z, keep_axes=0)
x, y, z = pt.unpack(p, shapes, keep_axes=0)

When pack is used in isolation it still makes sense for the user to specify keep_axes, as those are the ones that must match among different inputs being packed

But when unpack is used in isolation I think it makes less sense. You end up having to specify everything but the axis you want to operate on:

x = pt.tensor("x", shape=(8, 8, 10, 10))
pt.unpack(p, packed_shapes=[(2,), (2, 3)], keep_axes=[0, 2, 3])

If you never meant to pack these inputs (in which case it's nice that you can use the same keep_axes), it's not very ergonomic. It would be more intuitive to specify axis=1?

We don't even have a good default value if the input is >1D, as keep_axes=None could never produce such input. If you call the second example with the default keep_axes it raises.

But axis=0 if axis is None and keep_axes is None would be a reasonable default, and work with both workflows of unpack. I temporarily considered that default but decided to revert until we discuss having an explicit axis argument.

[jessegrabowski]

I think this makes sense. But if it's not already we're going to have to make sure everything is mandatory keyword only, because the op is already pretty complex.