Function diagonal

dask/array/creation.py:691–822

(a, offset=0, axis1=0, axis2=1)

@derived_from(np)
def diagonal(a, offset=0, axis1=0, axis2=1):
    name = "diagonal-" + tokenize(a, offset, axis1, axis2)

    if a.ndim < 2:
        # NumPy uses `diag` as we do here.
        raise ValueError("diag requires an array of at least two dimensions")

    def _axis_fmt(axis, name, ndim):
        if axis < 0:
            t = ndim + axis
            if t < 0:
                msg = "{}: axis {} is out of bounds for array of dimension {}"
                raise AxisError(msg.format(name, axis, ndim))
            axis = t
        return axis

    def pop_axes(chunks, axis1, axis2):
        chunks = list(chunks)
        chunks.pop(axis2)
        chunks.pop(axis1)
        return tuple(chunks)

    axis1 = _axis_fmt(axis1, "axis1", a.ndim)
    axis2 = _axis_fmt(axis2, "axis2", a.ndim)

    if axis1 == axis2:
        raise ValueError("axis1 and axis2 cannot be the same")

    a = asarray(a)
    k = offset
    if axis1 > axis2:
        axis1, axis2 = axis2, axis1
        k = -offset

    free_axes = set(range(a.ndim)) - {axis1, axis2}
    free_indices = list(product(*(range(a.numblocks[i]) for i in free_axes)))
    ndims_free = len(free_axes)

    # equation of diagonal: i = j - k
    kdiag_row_start = max(0, -k)
    kdiag_col_start = max(0, k)
    kdiag_row_stop = min(a.shape[axis1], a.shape[axis2] - k)
    len_kdiag = kdiag_row_stop - kdiag_row_start

    if len_kdiag <= 0:
        xp = np

        if is_cupy_type(a._meta):
            import cupy

            xp = cupy

        out_chunks = pop_axes(a.chunks, axis1, axis2) + ((0,),)
        dsk = {}
        for free_idx in free_indices:
            shape = tuple(
                out_chunks[axis][free_idx[axis]] for axis in range(ndims_free)
            )