Repository: dask/dask

Issue:
Dask Array's diag supporting diagonal selection
Would be nice if Dask Array's [`diag`]( http://dask.pydata.org/en/latest/array-api.html#dask.array.diag ) supported the `k` value to allow selection of the diagonal as NumPy's implementation [does]( https://docs.scipy.org/doc/numpy-1.13.0/reference/generated/numpy.diag.html ). It seems like Dask Array's API docs suggest `diag` does support `k`, but the code does not allow for this option.

https://github.com/dask/dask/blob/f015b5f0808420bdd23f8408f5822d2f5aff8ff1/dask/array/creation.py#L413


Discussion:
It might make sense to rewrite `diag` to use `diagonal` after PR ( https://github.com/dask/dask/pull/4431 ) is accepted. This would be a way to solve this problem and cutdown on code as well.
It's worth pointing out this feature is broken in both directions. You also can't make a dask 2-D array from a vector on k!=0 diagonal. Support for da.diagonal doesn't help at all in this case as da.diagonal only extracts diagonals from 2-D arrays. Here's a use case to reproduce the problem:

`import dask.array as da`
`v = da.random.uniform(size=100,chunks=50)`
`A = da.diag(v,k=1)`

which yields the following error:

`TypeError: diag() got an unexpected keyword argument 'k'`

At a minimum the documentation should be updated to reflect that you don't support the 'k' keyword argument. Currently the documentation says that it is supported.
http://docs.dask.org/en/latest/array-api.html#dask.array.diag
See https://github.com/dask/dask/issues/3265 or docstrings


On Sat, Mar 9, 2019 at 12:04 AM Mike Bauer <notifications@github.com> wrote:

> It's worth pointing out this feature is broken in both directions. You
> also can't make a dask 2-D array from a vector on k!=0 diagonal. Support
> for da.diagonal doesn't help at all in this case as da.diagonal only
> extracts diagonals from 2-D arrays. Here's a use case to reproduce the
> problem:
>
> `import dask.array as da
>
> v = da.random.uniform(size=100,chunks=50)
> A = da.diag(v,k=1)`
>
> which yields the following error:
>
> TypeError: diag() got an unexpected keyword argument 'k'
>
> At a minimum the documentation should be updated to reflect that you don't
> support the 'k' keyword argument. Currently the documentation says that it
> is supported.
> http://docs.dask.org/en/latest/array-api.html#dask.array.diag
>
> —
> You are receiving this because you are subscribed to this thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/2726#issuecomment-471156361>, or mute
> the thread
> <https://github.com/notifications/unsubscribe-auth/AASszK7vpzPRxhbs_tZ0MftX7jA7op32ks5vU2sJgaJpZM4PnhTO>
> .
>

I am working on a pull request on this using #4431 for the diagonal extraction case. WIP though, as I was surprised by the $(n+k)*(n+k)$ output of `np.diag(np.ones(n), k)`.
I also want to have a look into cleaning up the shared but different code in `da.triu`, `da.tril` and `da.diag`.
Cool, thanks for working on this. Please let us know if you need pointers or would like someone to take a look.

Resolving patch:
diff --git a/dask/array/creation.py b/dask/array/creation.py
--- a/dask/array/creation.py
+++ b/dask/array/creation.py
@@ -1,16 +1,15 @@
 import itertools
 from collections.abc import Sequence
-from functools import partial, reduce
+from functools import partial
 from itertools import product
 from numbers import Integral, Number
-from operator import getitem
 
 import numpy as np
 from tlz import sliding_window
 
 from ..base import tokenize
 from ..highlevelgraph import HighLevelGraph
-from ..utils import cached_cumsum, derived_from
+from ..utils import cached_cumsum, derived_from, is_cupy_type
 from . import chunk
 from .core import (
     Array,
@@ -565,8 +564,11 @@ def eye(N, chunks="auto", M=None, k=0, dtype=float):
 
 
 @derived_from(np)
-def diag(v):
-    name = "diag-" + tokenize(v)
+def diag(v, k=0):
+    if not isinstance(v, np.ndarray) and not isinstance(v, Array):
+        raise TypeError(f"v must be a dask array or numpy array, got {type(v)}")
+
+    name = "diag-" + tokenize(v, k)
 
     meta = meta_from_array(v, 2 if v.ndim == 1 else 1)
 
@@ -574,41 +576,51 @@ def diag(v):
         hasattr(v, "__array_function__") and not isinstance(v, Array)
     ):
         if v.ndim == 1:
-            chunks = ((v.shape[0],), (v.shape[0],))
-            dsk = {(name, 0, 0): (np.diag, v)}
+            m = abs(k)
+            chunks = ((v.shape[0] + m,), (v.shape[0] + m,))
+            dsk = {(name, 0, 0): (np.diag, v, k)}
         elif v.ndim == 2:
-            chunks = ((min(v.shape),),)
-            dsk = {(name, 0): (np.diag, v)}
+            kdiag_row_start = max(0, -k)
+            kdiag_row_stop = min(v.shape[0], v.shape[1] - k)
+            len_kdiag = kdiag_row_stop - kdiag_row_start
+            chunks = ((0,),) if len_kdiag <= 0 else ((len_kdiag,),)
+            dsk = {(name, 0): (np.diag, v, k)}
         else:
             raise ValueError("Array must be 1d or 2d only")
         return Array(dsk, name, chunks, meta=meta)
-    if not isinstance(v, Array):
-        raise TypeError(f"v must be a dask array or numpy array, got {type(v)}")
+
     if v.ndim != 1:
-        if v.chunks[0] == v.chunks[1]:
+        if v.ndim != 2:
+            raise ValueError("Array must be 1d or 2d only")
+        if k == 0 and v.chunks[0] == v.chunks[1]:
             dsk = {
                 (name, i): (np.diag, row[i]) for i, row in enumerate(v.__dask_keys__())
             }
             graph = HighLevelGraph.from_collections(name, dsk, dependencies=[v])
             return Array(graph, name, (v.chunks[0],), meta=meta)
         else:
-            raise NotImplementedError(
-                "Extracting diagonals from non-square chunked arrays"
-            )
-    chunks_1d = v.chunks[0]
-    blocks = v.__dask_keys__()
-    dsk = {}
-    for i, m in enumerate(chunks_1d):
-        for j, n in enumerate(chunks_1d):
-            key = (name, i, j)
-            if i == j:
-                dsk[key] = (np.diag, blocks[i])
-            else:
-                dsk[key] = (np.zeros, (m, n))
-                dsk[key] = (partial(np.zeros_like, shape=(m, n)), meta)
+            return diagonal(v, k)
+
+    if k == 0:
+        chunks_1d = v.chunks[0]
+        blocks = v.__dask_keys__()
+        dsk = {}
+        for i, m in enumerate(chunks_1d):
+            for j, n in enumerate(chunks_1d):
+                key = (name, i, j)
+                if i == j:
+                    dsk[key] = (np.diag, blocks[i])
+                else:
+                    dsk[key] = (np.zeros, (m, n))
+                    dsk[key] = (partial(np.zeros_like, shape=(m, n)), meta)
 
-    graph = HighLevelGraph.from_collections(name, dsk, dependencies=[v])
-    return Array(graph, name, (chunks_1d, chunks_1d), meta=meta)
+        graph = HighLevelGraph.from_collections(name, dsk, dependencies=[v])
+        return Array(graph, name, (chunks_1d, chunks_1d), meta=meta)
+
+    elif k > 0:
+        return pad(diag(v), [[0, k], [k, 0]], mode="constant")
+    elif k < 0:
+        return pad(diag(v), [[-k, 0], [0, -k]], mode="constant")
 
 
 @derived_from(np)
@@ -628,6 +640,12 @@ def _axis_fmt(axis, name, 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)
 
@@ -635,48 +653,105 @@ def _axis_fmt(axis, name, ndim):
         raise ValueError("axis1 and axis2 cannot be the same")
 
     a = asarray(a)
-
+    k = offset
     if axis1 > axis2:
         axis1, axis2 = axis2, axis1
-        offset = -offset
-
-    def _diag_len(dim1, dim2, offset):
-        return max(0, min(min(dim1, dim2), dim1 + offset, dim2 - offset))
-
-    diag_chunks = []
-    chunk_offsets = []
-    cum1 = cached_cumsum(a.chunks[axis1], initial_zero=True)[:-1]
-    cum2 = cached_cumsum(a.chunks[axis2], initial_zero=True)[:-1]
-    for co1, c1 in zip(cum1, a.chunks[axis1]):
-        chunk_offsets.append([])
-        for co2, c2 in zip(cum2, a.chunks[axis2]):
-            k = offset + co1 - co2
-            diag_chunks.append(_diag_len(c1, c2, k))
-            chunk_offsets[-1].append(k)
+        k = -offset
 
-    dsk = {}
-    idx_set = set(range(a.ndim)) - {axis1, axis2}
-    n1 = len(a.chunks[axis1])
-    n2 = len(a.chunks[axis2])
-    for idx in product(*(range(len(a.chunks[i])) for i in idx_set)):
-        for i, (i1, i2) in enumerate(product(range(n1), range(n2))):
-            tsk = reduce(getitem, idx[:axis1], a.__dask_keys__())[i1]
-            tsk = reduce(getitem, idx[axis1 : axis2 - 1], tsk)[i2]
-            tsk = reduce(getitem, idx[axis2 - 1 :], tsk)
-            k = chunk_offsets[i1][i2]
-            dsk[(name,) + idx + (i,)] = (np.diagonal, tsk, k, axis1, axis2)
-
-    left_shape = tuple(a.shape[i] for i in idx_set)
-    right_shape = (_diag_len(a.shape[axis1], a.shape[axis2], offset),)
-    shape = left_shape + right_shape
-
-    left_chunks = tuple(a.chunks[i] for i in idx_set)
-    right_shape = (tuple(diag_chunks),)
-    chunks = left_chunks + right_shape
+    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 = dict()
+        for free_idx in free_indices:
+            shape = tuple(
+                out_chunks[axis][free_idx[axis]] for axis in range(ndims_free)
+            )
+            dsk[(name,) + free_idx + (0,)] = (
+                partial(xp.empty, dtype=a.dtype),
+                shape + (0,),
+            )
+
+        meta = meta_from_array(a, ndims_free + 1)
+        return Array(dsk, name, out_chunks, meta=meta)
+
+    # compute row index ranges for chunks along axis1:
+    row_stops_ = np.cumsum(a.chunks[axis1])
+    row_starts = np.roll(row_stops_, 1)
+    row_starts[0] = 0
+
+    # compute column index ranges for chunks along axis2:
+    col_stops_ = np.cumsum(a.chunks[axis2])
+    col_starts = np.roll(col_stops_, 1)
+    col_starts[0] = 0
+
+    # locate first chunk containing diagonal:
+    row_blockid = np.arange(a.numblocks[axis1])
+    col_blockid = np.arange(a.numblocks[axis2])
+
+    row_filter = (row_starts <= kdiag_row_start) & (kdiag_row_start < row_stops_)
+    col_filter = (col_starts <= kdiag_col_start) & (kdiag_col_start < col_stops_)
+    (I,) = row_blockid[row_filter]
+    (J,) = col_blockid[col_filter]
+
+    # follow k-diagonal through chunks while constructing dask graph:
+    dsk = dict()
+    i = 0
+    kdiag_chunks = ()
+    while kdiag_row_start < a.shape[axis1] and kdiag_col_start < a.shape[axis2]:
+        # localize block info:
+        nrows, ncols = a.chunks[axis1][I], a.chunks[axis2][J]
+        kdiag_row_start -= row_starts[I]
+        kdiag_col_start -= col_starts[J]
+        k = -kdiag_row_start if kdiag_row_start > 0 else kdiag_col_start
+        kdiag_row_end = min(nrows, ncols - k)
+        kdiag_len = kdiag_row_end - kdiag_row_start
+
+        # increment dask graph:
+        for free_idx in free_indices:
+            input_idx = (
+                free_idx[:axis1]
+                + (I,)
+                + free_idx[axis1 : axis2 - 1]
+                + (J,)
+                + free_idx[axis2 - 1 :]
+            )
+            output_idx = free_idx + (i,)
+            dsk[(name,) + output_idx] = (
+                np.diagonal,
+                (a.name,) + input_idx,
+                k,
+                axis1,
+                axis2,
+            )
+
+        kdiag_chunks += (kdiag_len,)
+        # prepare for next iteration:
+        i += 1
+        kdiag_row_start = kdiag_row_end + row_starts[I]
+        kdiag_col_start = min(ncols, nrows + k) + col_starts[J]
+        I = I + 1 if kdiag_row_start == row_stops_[I] else I
+        J = J + 1 if kdiag_col_start == col_stops_[J] else J
 
+    out_chunks = pop_axes(a.chunks, axis1, axis2) + (kdiag_chunks,)
     graph = HighLevelGraph.from_collections(name, dsk, dependencies=[a])
-    meta = meta_from_array(a, len(shape))
-    return Array(graph, name, shape=shape, chunks=chunks, meta=meta)
+    meta = meta_from_array(a, ndims_free + 1)
+    return Array(graph, name, out_chunks, meta=meta)
 
 
 @derived_from(np)


Test patch:
diff --git a/dask/array/tests/test_creation.py b/dask/array/tests/test_creation.py
--- a/dask/array/tests/test_creation.py
+++ b/dask/array/tests/test_creation.py
@@ -407,32 +407,70 @@ def test_eye():
         assert 4 < x.npartitions < 32
 
 
-def test_diag():
+@pytest.mark.parametrize("k", [0, 3, -3, 8])
+def test_diag_bad_input(k):
+    # when input numpy array is neither 1d nor 2d:
+    v = np.arange(2 * 3 * 4).reshape((2, 3, 4))
+    with pytest.raises(ValueError, match="Array must be 1d or 2d only"):
+        da.diag(v, k)
+
+    # when input dask array is neither 1d nor 2d:
+    v = da.arange(2 * 3 * 4).reshape((2, 3, 4))
+    with pytest.raises(ValueError, match="Array must be 1d or 2d only"):
+        da.diag(v, k)
+
+    # when input is not an array:
+    v = 1
+    with pytest.raises(TypeError, match="v must be a dask array or numpy array"):
+        da.diag(v, k)
+
+
+@pytest.mark.parametrize("k", [0, 3, -3, 8])
+def test_diag_2d_array_creation(k):
+    # when input 1d-array is a numpy array:
     v = np.arange(11)
-    assert_eq(da.diag(v), np.diag(v))
+    assert_eq(da.diag(v, k), np.diag(v, k))
 
+    # when input 1d-array is a dask array:
     v = da.arange(11, chunks=3)
-    darr = da.diag(v)
-    nparr = np.diag(v)
+    darr = da.diag(v, k)
+    nparr = np.diag(v, k)
     assert_eq(darr, nparr)
-    assert sorted(da.diag(v).dask) == sorted(da.diag(v).dask)
+    assert sorted(da.diag(v, k).dask) == sorted(da.diag(v, k).dask)
 
     v = v + v + 3
-    darr = da.diag(v)
-    nparr = np.diag(v)
+    darr = da.diag(v, k)
+    nparr = np.diag(v, k)
     assert_eq(darr, nparr)
 
     v = da.arange(11, chunks=11)
-    darr = da.diag(v)
-    nparr = np.diag(v)
+    darr = da.diag(v, k)
+    nparr = np.diag(v, k)
     assert_eq(darr, nparr)
-    assert sorted(da.diag(v).dask) == sorted(da.diag(v).dask)
+    assert sorted(da.diag(v, k).dask) == sorted(da.diag(v, k).dask)
 
-    x = np.arange(64).reshape((8, 8))
-    assert_eq(da.diag(x), np.diag(x))
 
+@pytest.mark.parametrize("k", [0, 3, -3, 8])
+def test_diag_extraction(k):
+    # when input 2d-array is a square numpy array:
+    x = np.arange(64).reshape((8, 8))
+    assert_eq(da.diag(x, k), np.diag(x, k))
+    # when input 2d-array is a square dask array:
     d = da.from_array(x, chunks=(4, 4))
-    assert_eq(da.diag(d), np.diag(x))
+    assert_eq(da.diag(d, k), np.diag(x, k))
+    # heterogeneous chunks:
+    d = da.from_array(x, chunks=((3, 2, 3), (4, 1, 2, 1)))
+    assert_eq(da.diag(d, k), np.diag(x, k))
+
+    # when input 2d-array is a rectangular numpy array:
+    y = np.arange(5 * 8).reshape((5, 8))
+    assert_eq(da.diag(y, k), np.diag(y, k))
+    # when input 2d-array is a rectangular dask array:
+    d = da.from_array(y, chunks=(4, 4))
+    assert_eq(da.diag(d, k), np.diag(y, k))
+    # heterogeneous chunks:
+    d = da.from_array(y, chunks=((3, 2), (4, 1, 2, 1)))
+    assert_eq(da.diag(d, k), np.diag(y, k))
 
 
 def test_diagonal():