Function squeeze

dask/array/routines.py:1991–2013

(a, axis=None)

@derived_from(np)
def squeeze(a, axis=None):
    if axis is None:
        axis = tuple(i for i, d in enumerate(a.shape) if d == 1)
    elif not isinstance(axis, tuple):
        axis = (axis,)

    if any(a.shape[i] != 1 for i in axis):
        raise ValueError("cannot squeeze axis with size other than one")

    axis = validate_axis(axis, a.ndim)

    sl = tuple(0 if i in axis else slice(None) for i, s in enumerate(a.shape))

    # Return 0d Dask Array if all axes are squeezed,
    # to be consistent with NumPy. Ref: https://github.com/dask/dask/issues/9183#issuecomment-1155626619
    if all(s == 0 for s in sl) and all(s == 1 for s in a.shape):
        return a.map_blocks(
            np.squeeze, meta=a._meta, drop_axis=tuple(range(len(a.shape)))
        )

    a = a[sl]

    return a


@derived_from(np)