def build_argmax_or_argmin_with_axis_impl(a, axis, flatten_impl):
    """
    Given a function that implements the logic for handling a flattened
    array, return the implementation function.
    """
    check_is_integer(axis, "axis")
    retty = types.intp

    tuple_buffer = tuple(range(a.ndim))

    def impl(a, axis=None):
        if axis < 0:
            axis = a.ndim + axis

        if axis < 0 or axis >= a.ndim:
            raise ValueError("axis is out of bounds")

        # Short circuit 1-dimensional arrays:
        if a.ndim == 1:
            return flatten_impl(a)

        # Make chosen axis the last axis:
        tmp = tuple_buffer
        for i in range(axis, a.ndim - 1):
            tmp = tuple_setitem(tmp, i, i + 1)
        transpose_index = tuple_setitem(tmp, a.ndim - 1, axis)
        transposed_arr = a.transpose(transpose_index)

        # Flatten along that axis; since we've transposed, we can just get
        # batches off the overall flattened array.
        m = transposed_arr.shape[-1]
        raveled = transposed_arr.ravel()
        assert raveled.size == a.size
        assert transposed_arr.size % m == 0
        out = np.empty(transposed_arr.size // m, retty)
        for i in range(out.size):
            out[i] = flatten_impl(raveled[i * m:(i + 1) * m])

        # Reshape based on axis we didn't flatten over:
        return out.reshape(transposed_arr.shape[:-1])

    return impl

    Given a function that implements the logic for handling a flattened
    array, return the implementation function.
    
numba/np/old_arraymath.py