Repository: dask/dask
File: dask/array/creation.py (Python)

Code before:
    if v.ndim != 1:
        if v.ndim != 2:
            raise ValueError("Array must be 1d or 2d only")
        if k == 0 and v.chunks[0] == v.chunks[1]:
            tasks = [
                Task((name, i), np.diag, TaskRef(row[i]))
                for i, row in enumerate(v.__dask_keys__())
            ]
            dsk = {t.key: t for t in tasks}
            graph = HighLevelGraph.from_collections(name, dsk, dependencies=[v])
            return Array(graph, name, (v.chunks[0],), meta=meta)
        else:
            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] = Task(key, np.diag, TaskRef(blocks[i]))
                else:
                    # dsk[key] = Task(key, np.zeros, (m, n))
                    dsk[key] = Task(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)

    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)
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

Review comment (suggestion):
FYI No need for partials (but it shouldn't be harmful). `Task` accepts kwargs as well
```suggestion
                    dsk[key] = Task(key, np.zeros_like, meta, shape=(m, n))
```

Code after:
    if v.ndim != 1:
        if v.ndim != 2:
            raise ValueError("Array must be 1d or 2d only")
        if k == 0 and v.chunks[0] == v.chunks[1]:
            tasks = [
                Task((name, i), np.diag, TaskRef(row[i]))
                for i, row in enumerate(v.__dask_keys__())
            ]
            dsk = {t.key: t for t in tasks}
            graph = HighLevelGraph.from_collections(name, dsk, dependencies=[v])
            return Array(graph, name, (v.chunks[0],), meta=meta)
        else:
            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] = Task(key, np.diag, TaskRef(blocks[i]))
                else:
                    dsk[key] = Task(key, np.zeros_like, meta, shape=(m, n))

        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)
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