Segmenting numpy arrays with as_strided

Based on DanielF's comment and his answer here, I implemented my function like this:

def segment(arr, axis, new_len, step=1, new_axis=None, return_view=False):    """ Segment an array along some axis.    Parameters    ----------    arr : array-like        The input array.    axis : int        The axis along which to segment.    new_len : int        The length of each segment.    step : int, default 1        The offset between the start of each segment.    new_axis : int, optional        The position where the newly created axis is to be inserted. By        default, the axis will be added at the end of the array.    return_view : bool, default False        If True, return a view of the segmented array instead of a copy.    Returns    -------    arr_seg : array-like        The segmented array.    """    old_shape = np.array(arr.shape)    assert new_len <= old_shape[axis],  \        "new_len is bigger than input array in axis"    seg_shape = old_shape.copy()    seg_shape[axis] = new_len    steps = np.ones_like(old_shape)    if step:        step = np.array(step, ndmin = 1)        assert step > 0, "Only positive steps allowed"        steps[axis] = step    arr_strides = np.array(arr.strides)    shape = tuple((old_shape - seg_shape) // steps + 1) + tuple(seg_shape)    strides = tuple(arr_strides * steps) + tuple(arr_strides)    arr_seg = np.squeeze(        as_strided(arr, shape = shape, strides = strides))    # squeeze will move the segmented axis to the first position    arr_seg = np.moveaxis(arr_seg, 0, axis)    # the new axis comes right after    if new_axis is not None:        arr_seg = np.moveaxis(arr_seg, axis+1, new_axis)    else:        arr_seg = np.moveaxis(arr_seg, axis+1, -1)    if return_view:        return arr_seg    else:        return arr_seg.copy()

This works well for my case of one-dimensional segments, however, I'd recommend anyone looking for a way that works for segments of arbitrary dimensionality to check out the code in the linked answer.