scipy.signal.welch performance improvements with array API

[tylerjereddy]

Following up on some work I was doing at scipy/scipy#18286 and tylerjereddy/scipy#70 , where scipy.signal.welch() was ~10x faster with CuPy vs. NumPy for 7e7 points, I took a look at what would happen if I tried to pass in PyKokkos views instead of NumPy/CuPy arrays.

To start with, the tests will fail because we don't have: __array_namespace__

While that may not be too hard to support, this example of GPU speedup in SciPy will also require us to provide a complete FFT implementation, which seems like a heavy lift unless the Kokkos ecosystem already provides it somewhere: https://data-apis.org/array-api/2022.12/extensions/fourier_transform_functions.html

Also, the testsuite had many other complaints about trying to use PyKokkos as a backend, but that's a good start, and already a decent chunk of work if you want to compete with the other libs over there. Some other missing things were arange, empty, complex128...

If you want to try it for yourself you can use that branch of SciPy with this mod:

diff --git a/scipy/_lib/_testutils.py b/scipy/_lib/_testutils.py
index 61762c590..df185c292 100644
--- a/scipy/_lib/_testutils.py
+++ b/scipy/_lib/_testutils.py
@@ -288,6 +288,8 @@ def _import_xp():
             import cupy as xp
         elif backend == "pytorch":
             import torch as xp
+        elif backend == "pykokkos":
+            import pykokkos as xp
         else:
             raise ValueError(f"ARR_TST_BACKEND {backend} not recognized.")
     except KeyError:

Then run i.e., : ARR_TST_BACKEND=pykokkos python dev.py test -j 12 -t scipy/signal/tests/test_spectral.py

Development is moving fast though... that was at hash e5287dd on my branch of SciPy.

Also, note the importance of providing as_strided or similar: scipy/scipy#18286 (comment)