Repository: numpy/numpy
File: numpy/lib/_function_base_impl.py (Python)

Code before:
            nkeep = len(keep)

            def reshape_arr(a):
                # swap axis that should not be reduced to front
                for i, s in enumerate(sorted(keep)):
                    a = a.swapaxes(i, s)
                # merge reduced axis
                return a.reshape(a.shape[:nkeep] + (-1,))

            a = reshape_arr(a)

            weights = kwargs.get("weights")
            if weights is not None:
                weights = np.asanyarray(weights)
                kwargs["weights"] = reshape_arr(weights)

            kwargs['axis'] = -1
    elif keepdims and out is not None:
        index_out = (0, ) * nd
        kwargs['out'] = out[(Ellipsis, ) + index_out]

    r = func(a, **kwargs)

    if out is not None:
        return out

    if keepdims:
        if axis is None:
            index_r = (np.newaxis, ) * nd
        else:
            index_r = tuple(
                np.newaxis if i in axis else slice(None)
                for i in range(nd))
        r = r[(Ellipsis, ) + index_r]

    return r


def _median_dispatcher(
        a, axis=None, out=None, overwrite_input=None, keepdims=None):
    return (a, out)


@array_function_dispatch(_median_dispatcher)
def median(a, axis=None, out=None, overwrite_input=False, keepdims=False):
    """
    Compute the median along the specified axis.

    Returns the median of the array elements.

    Parameters

Review comment (refactor):
Hehe, for once einops-like syntax would be nice for us also.

I do wonder if it isn't nicer to just duplicate things by using `moveaxis`.
```
a = np.moveaxis(a, keep, range(nkeep))
a = a.reshape(a.shape[:nkeep] + (-1,)) 
```
And then just repeat, because it's just two lines?  The `asanyarray(weights)` below should not be necessary here, I think.  We need additional preparation for the `weights` already (because they have funky logic for when the weights match the number of reduction axes)

Code after:
    axis = kwargs.get('axis')
    out = kwargs.get('out')

    if keepdims is np._NoValue:
        keepdims = False

    nd = a.ndim
    if axis is not None:
        axis = _nx.normalize_axis_tuple(axis, nd)

        if keepdims and out is not None:
            index_out = tuple(
                0 if i in axis else slice(None) for i in range(nd))
            kwargs['out'] = out[(Ellipsis, ) + index_out]

        if len(axis) == 1:
            kwargs['axis'] = axis[0]
        else:
            keep = sorted(set(range(nd)) - set(axis))
            nkeep = len(keep)

            def reshape_arr(a):
                # move axis that should not be reduced to front
                a = np.moveaxis(a, keep, range(nkeep))
                # merge reduced axis
                return a.reshape(a.shape[:nkeep] + (-1,))

            a = reshape_arr(a)

            weights = kwargs.get("weights")
            if weights is not None:
                kwargs["weights"] = reshape_arr(weights)

            kwargs['axis'] = -1
    elif keepdims and out is not None:
        index_out = (0, ) * nd
        kwargs['out'] = out[(Ellipsis, ) + index_out]

    r = func(a, **kwargs)

    if out is not None:
        return out

    if keepdims:
        if axis is None:
            index_r = (np.newaxis, ) * nd
        else:
            index_r = tuple(
                np.newaxis if i in axis else slice(None)
                for i in range(nd))
        r = r[(Ellipsis, ) + index_r]