ENH: add `split`, `swapaxes`

[ev-br]

A discussion and a crude implementation of part of np.split and np.swapaxes is at scipy/scipy#22122 (comment)

[kgryte]

@ev-br The specification has moveaxis: https://data-apis.org/array-api/latest/API_specification/generated/array_api.moveaxis.html.

[ev-br]

@kgryte yes. swapaxes may be implemented either with moveaxis or via permute_dims as in https://github.com/scipy/scipy/pull/22122/files#r1916726566. Or may stay a private utility in scipy/signal/_signaltools.

[kgryte]

I suppose my question is why is swapaxes necessary when there is moveaxis?

[ev-br]

It's pure convenience indeed, not a necessity. Esp when converting legacy code which predates moveaxis in numpy. A single call to swapaxes vs two calls to moveaxis, and since swapaxis semantics is confusing, so is replicating it with moveaxis. New code is indeed way better off using moveaxis.

[lucascolley]

two calls to moveaxis

FWIW it is just one call, swapaxes(x, y, z) == moveaxis(x, (y, z), (z, y)). I'm indifferent to adding that alias to array-api-extra, or just using xp.moveaxis directly in SciPy, but both are better than leaving a private helper in the codebase for it.

[lucascolley]

swapaxes(x, y, z) == moveaxis(x, (y, z), (z, y))

It turns out that this isn't quite true. xp.moveaxis(x, (axis, -1), (-1, axis)) fails for axis == -1 (modulo x.ndim). So anyone converting from np.swapaxes still needs to sanitise input for the case of y == z modulo x.ndim.

I can see the argument for wanting to abstract this sanitisation out into a separate helper when converting code from NumPy.

scipy/scipy#20639 (comment)

[lucascolley]

An updated implementation would look something like

    ndim = a.ndim
    if (a1 := axis1 % ndim) == (a2 := axis2 % ndim):
        return a
    return xp.moveaxis(a, (a1, a2), (a2, a1))

or without the walruses

    ndim = a.ndim
    axis1 = axis1 % ndim
    axis2 = axis2 % ndim
    if axis1 == axis2:
        return a
    return xp.moveaxis(a, (axis1, axis2), (axis2, axis1))