✔️ Linting Passed

All linting checks passed. Your pull request is in excellent shape! ☀️

_{Generated for commit: 7c237dc. Link to the linter CI: here}

Maybe we should go back to using .flat for numpy and a for loop for non-numpy, similarly to what it was before #29433?

I have a PoC in a branch: https://github.com/lesteve/scikit-learn/blob/4a270381c2c21c715f58923893df87b8c503af0f/sklearn/utils/_array_api.py#L530-L575

The main arguments I have for using .flat/loop (vs reshape + use return value):

• avoid a copy in the numpy case for a non C-contiguous array. This also avoids a copy for other namespaces but needs to do a for loop.
• slightly less tricky code maybe? I guess this avoids the caveat that .reshape can be a copy in non C-contiguous arrays.
• slightly harder to misuse. It's kind of easy to think that _fill_or_add_diagonal is in-place like np.fill_diagonal and not use the return value. This will mostly work in tests because array will be C-contiguous but will be broken in non C-contiguous cases as we discovered.

Thanks @lesteve ! You make some good points, I'm pretty convinced, especially by your last point. I had a look at your implementation, which is clever but I thought it was a bit confusing with the assignment_function.

I created another implementation option and tested the performance your implementation, the original implementation and my new implemenation. Here are the functions in a gist. Summary of the 6 functions tested:

• fill_reshape and add_reshape -> using the original _fill_or_add_to_diagonal, which is named _fill_or_add_to_diagonal_reshape in the gist, as it uses reshape. I checked with both add_value as True and False.
• fill_loop and add_loop -> my implementation that updates value (with diagonal() + value when add_value=True so we always assign with arr[i,i] = value[i]). I checked with both add_value` as True and False
• fill_helper and add_helper your implemention, using the _fill_diagonal_helper

With array (xp.zeros) of size 1e4, 1e4 (note I couldn't make the array much bigger as it resulted in maxing out all the memory in my colab session):

Benchmarking with array shape (10000, 10000)

Timing Results (seconds):
------------------------------------------------------------
Function             NumPy           CuPy           
------------------------------------------------------------
fill_reshape               0.000449       0.001825
add_reshape                0.000523       0.001568
fill_loop                  0.000853       0.157981
add_loop                   0.000919       0.106993
fill_helper                0.000714       0.102225
add_helper                 0.001321       0.178947

With transposed array (xp.zeros) of size 1e4, 1e4:

Benchmarking with array shape (10000, 10000)

Timing Results (seconds):
------------------------------------------------------------
Function             NumPy           CuPy           
------------------------------------------------------------
fill_reshape               1.268526       0.017735
add_reshape                1.268547       0.017209
fill_loop                  0.000769       0.106146
add_loop                   0.001124       0.101842
fill_helper                0.001136       0.101706
add_helper                 0.001776       0.180543

I also tried to measure the difference in peak memory usage with numpy with array and transposed array with memory_profiler.

C-contiguous array:

Profiling normal array:
Array info:
  Strides: (80000, 8)
  C-contiguous: True
  F-contiguous: False

Memory profile:
  Initial memory: 1010.2 MB
  Peak memory:    1010.6 MB
  Memory increase: 0.4 MB
  Final memory:   1010.6 MB

Transposed F-contiguous array

Profiling transposed array:
Array info:
  Strides: (8, 80000)
  C-contiguous: False
  F-contiguous: True

Memory profile:
  Initial memory: 1010.6 MB
  Peak memory:    1773.6 MB
  Memory increase: 763.0 MB
  Final memory:   1773.6 MB

Performance:

• Transposed array slows down the 'reshape' implementation but no consistent change to other 2 implementations that use loops
• The 'reshape' implementation is much faster for non transposed arrays but peak memory is higher
• My implementation is faster than yours when adding value but yours is faster for fill value

wrap

• Currently we never use wrap=True.
• It would be difficult to implement wrap=True using diagonal() + value as diagonal does not have a wrap parameter

Maybe we could have 2 separate functions, with fill_diagonal supporting wrap (just because numpy does) and add_to_diagonal not supporting wrap ?

Summary:

reshape is much faster for C-contiguous arrays, which is a shame, but I think I agree with your points (note that we avoid copy also for torch, which also returned a copy). Let me push a suggestion using my implementation for add and yours for fill.

I have now gone down this rabbit hole way too far 😅

Test failures are because array-api-strict doesn't have diagonal which is technically in the linalg section.
The array api compat has:

which does not work for array api strict...?

Edit: Ignore this, I did not realise there was a linalg extension 😬