increased performance of k-diagonal extraction in da.diag() and da.diagonal()

[ParticularMiner]

• Closes Dask Array's diag supporting diagonal selection #2726 (and maybe [WIP] Add k arg to diag #5683)
• Tests added / passed
• Passes pre-commit run --all-files

Perhaps you might find this useful. It mirrors the dask-grblas implementation, which supports rectangular chunks.

It follows the straight path of the k-diagonal through the rectangular chunks of the input matrix , constructing the dask graph along the way. So chunks untouched by the diagonal end up not being part of the final dask graph, reducing the algorithmic complexity of diagonal(A, offset, axis1, axis2) from

O(M N_axis1 N_axis2)       to        O(M max(1, N_k)) ,

where N_axis1 (N_axis2) is the number of chunks along axis1 (axis2) of array A; M is the total number of chunks of A after axis1 and axis2 have been removed; while N_k is the number of chunks touched by the k-diagonal after all axes, except axis1 and axis2, have been removed.

Relevant test-units have been modified to reflect this change.

Critique is welcome.

NB: It might be worth comparing this with #5683 (which I only discovered after pushing this commit).

[GPUtester]

Can one of the admins verify this patch?

[quasiben]

add to allowlist

[ParticularMiner]

FYI, I have extended this PR with @TAdeJong's simple padding solution for diag(v, k) when v is 1d and k != 0.

[TAdeJong]

At the time, #5683 failed on the issue described in #5661 with test_corrcoef() in dask\array\tests\test_routines.py (which was semi-unrelated). As of master today, #5661 still seems valid, so it might be interesting in comparing to see if this does not trigger #5661 and if so: why not.

[TAdeJong]

For completeness:
I am the original author of #5683 and happy to let this PR supersede it 👍
As per discussion there, #5661 is essentially independent of this PR. (Which makes this all the better PR to merge of the 2).

[TAdeJong]

Other than thinking about not duplicating code from da.diagonal as mentioned above, this PR seems in all ways better than #5683 to me.
I think (admittedly from the sidelines, it has been some time since I was this deep in the dask code base) that if @ParticularMiner already had a good reason to not call da.diagonal this can be merged, #5683 can be closed/rejected and #8703 can be handled separately, otherwise it might be worth it to incorporate #8703 here, call da.diagonal and review and merge this as one PR.

[ParticularMiner]

@TAdeJong

Good question. You're right about code-duplication issues in that code-section. I would do so but ...

I feel it would be less performant to call diagonal()'s code from diag(). The reason being that diagonal() converts numpy arrays to dask arrays before constructing the dask graph which later extracts the k-diagonal, while diag() takes a more direct route to the graph without such any such conversion.

On the other hand, performance here is up for discussion as very large numpy arrays could benefit from being converted to dask arrays prior to graph construction. As a rule, however, I try to avoid creating very large numpy arrays in the first place.

I don't feel strongly though about performance in this particular case, so if you have other arguments in favor of avoiding code duplication, feel free to share!

[ParticularMiner]

Btw, many thanks for your constructive critique and support of this PR @TAdeJong .

[ParticularMiner]

@TAdeJong , @pavithraes , @ian-r-rose

RE: code duplication

Upon further reflection, I agree with @TAdeJong's suggestion to reduce code-duplication.

To do this, I generalized this k-diagonal extraction algorithm to higher dimensions and transferred the code into diagonal(), so that diag() now calls diagonal() where applicable. Now both diag() and diagonal() are independent of #5661.

See the opening comment above of this PR for a note on algorithmic complexity/performance.

[jcrist]

Thanks @ParticularMiner , overall this looks good to me. I left one test comment, but the rest looks good. While you're looking at this code, can I ask that you look through the tests for da.diagonal as well to see if they offer full coverage for the different parameters we'd care about here (dimensions of input array, chunking variation, axis arguments, etc...)? It'd be good to ensure we have full coverage here while you're still thinking about this code.

[pavithraes]

Thanks for the update @ParticularMiner!

@jcrist Did the updates address your comments?

[jcrist]

🚀

[jakirkham]

Thanks all! 😄