Tracker: array types (CuPy, PyTorch & co) and array API standard support

[rgommers]

See gh-18286 for the proposed design changes, and gh-18668 for the main PR that laid the foundations (it supported all of scipy.cluster, added a CI job that used both pytorch-cpu and numpy.array_api, and documented the design patterns).

This issue is to track the current status of support across submodules, and other larger tasks/TODOs. Submodules:

• scipy.cluster
• scipy.constants ENH: constants: add array api support #20593
• scipy.datasets ENH: datasets: array API standard support #20594 (no changes needed)
• scipy.fft
• scipy.fftpack # Don't do, fftpack is legacy
• scipy.integrate
• scipy.interpolate
• scipy.io
• scipy.linalg (in progress, see ENH: linalg: array library interoperability #19068)
• scipy.misc # Don't do, will be removed
• scipy.ndimage
• scipy.odr
• scipy.optimize
• scipy.signal ENH: signal: add array API support #20678
• scipy.sparse # TBD on whether to update functions in that namespace; don't touch the data structures
• scipy.sparse.csgraph
• scipy.sparse.linalg
• scipy.spatial
• scipy.special (initial set of functions done by @mdhaber)
• scipy.stats (see ENH: stats: add array API-support #20544 - in progress by @mdhaber, PRs welcome from others)

Desirable for later:

• Once we're more confident that everything is in good shape, see if we should start emitting warnings for input types that may see a change in behavior (see this comment).
• Write a policy for cases where one wants to change existing pure Python code with array API support to compiled code (for performance improvements with numpy). See this comment).
• Figure out if and how to deal with mixing multiple array types coming from the same library (see this comment).
• Once JAX and/or Dask are supported in the array-api-compat package, add support for them in SciPy (both can be tested in CI too). (JAX in progress, see ENH: array types: add JAX support #20085)
• CuPy isn't testable in (free) CI easily, but it could be done (e.g., with https://cirun.io/). Shouldn't be done before all this becomes stable/supported.

[izaid]

I have a question that probably falls under "desirable for later". Regarding scipy.sparse, is the intention to update the various sparse array / matrix classes to use the array api? Or has there been some discussion about making use of the sparse classes that already exist in CuPy, PyTorch, etc?

[rgommers]

Good question @izaid. For sparse it's mainly the functionality on top of the data structures themselves that can support multiple array types - I've split off sparse.csgraph and sparse.linalg in the list above. Major changes to the sparse API is a separate topic; that's not in the cards I think. There's some desire from some folks to deprecate *_matrix in favor of *_array, whether that happens and on what timeline is still TBD and waiting for a concrete proposal. The *_array API should see some additions, but I can't see it become array API standard compatible (that'd be pydata/sparse).

[ivirshup]

I've opened #18915 to discuss updating the sparse array class to adopt the array api as a provider.

I am curious what functionality in sparse and its submodules could be targeted for becoming array API consumers given the broad reliance on native code. Has anyone looked into this yet?

[rgommers]

Thanks @ivirshup! I did think about that quite a bit already. I'm a bit swamped right this minute, but will to reply on gh-18915 within 1-2 days.

[lucascolley]

We could add that partial coverages of signal and linalg are in progress by Tyler and me respectively.

[lucascolley]

fft can be ticked off now!

[lucascolley]

Write a policy for cases where one wants to change existing pure Python code with array API support to compiled code (for performance improvements with numpy)

I think this should be relatively straightforward - we just need to maintain two code branches, one which is pure python and one which is in whichever language we want for NumPy arrays.

In fact, I think it would be useful to ensure that no pure python code (including submodules which do not yet have array API support) is removed from now on. For example, if a Python file gets Cythonized, the pure Python file should be kept (even if it is 'dead code' at the time).

[rgommers]

the pure Python file should be kept (even if it is 'dead code' at the time).

That doesn't sound like a reasonable restriction to me. We have lots of compiled code around, and there isn't a real need to preemptively make it harder to accelerate code.

I think this should be relatively straightforward - we just need to maintain two code branches, one which is pure python and one which is in whichever language we want for NumPy arrays.

This I agree with. The point being that we want to avoid a regression in functionality - once something works with a GPU library for example, we can't break that. Which means keeping pure Python code around (it could be rewritten of course if that'd make it better/faster - a lot of our code is really old and suboptimal).

[lucascolley]

That doesn't sound like a reasonable restriction to me. We have lots of compiled code around, and there isn't a real need to preemptively make it harder to accelerate code.

In that case I would suggest that any new translations at least add a comment into the code to state that they are translations from pure Python. That should avoid a scenario where we mistakenly 'dismiss' a module as compiled code rather than going into the git history and restoring (an array-agnostic version of) the pure Python implementation.

[rgommers]

Sure, a code comment won't hurt.

[dschmitz89]

One general question: is there a best practice to install the additional dependencies in the scipy dev environment? The thought of adding pytorch, cupy and potentially jax to my conda environment is a little scary and likely to mess everything up. Might be worth to document also.

[andyfaff]

This is why conda environments are good, you can just remove them if you mess them up. You just have to remember to not install things into the base environment.

[dschmitz89]

True, but my issue is that in the past I sometimes had to completely nuke and recreate the scipy development environment only for scipy itself. Getting the scipy development environment working is also a big issue for many newcomers.

These additional heavy dependencies will make the whole dev environment even more prone to build problems. That's why I asked for a guideline if there is one.

[lucascolley]

One Conda environment with every package is... challenging. For me, on Ubuntu, installing the correct Nvidia drivers and compatible CuPy/PyTorch/JAX versions in one environment is something I haven't managed to do yet.

See for example google/jax#18032 (comment), which shows that you'll want CUDA 12.1. But I still haven't been able to figure out which parts are supposed to be downloaded from the Nvidia website, and which things should come straight from Conda.

This is further complicated by the fact that I use a (potentially different) driver version for my system, and the process to uninstall/switch versions is not clear.

One thing that I have found quite easy is to have separate environments (SciPy + JAX, SciPy + PyTorch etc.). That wasn't too difficult to get working, but it's not ideal.

[lucascolley]

With that said, I agree that this needs documentation and a streamlined process. There are too many moving parts at the minute though IMO, especially on GPU, and covering GPU drivers across different platforms is a big task.

[rgommers]

What I do is this:

1. Edit environment.yml to add the desired dependencies for any custom testing, and append something descriptive to the default scipy-devenvironment name (e.g.,scipy-dev-cupy`)
2. Run `mamba env create -f environment.yml
3. Undo changes to environment.yml

Results:

$ mamba env list | rg scipy
scipy-dev             *  /home/rgommers/mambaforge/envs/scipy-dev
scipy-dev-icx            /home/rgommers/mambaforge/envs/scipy-dev-icx
scipy-dev-jax            /home/rgommers/mambaforge/envs/scipy-dev-jax
scipy-dev-sparse         /home/rgommers/mambaforge/envs/scipy-dev-sparse
....

Everything in a single env may work, but is indeed likely to be fragile. No reason not to have a bunch of separate envs.

In a similar vein: I also have a couple of repo clones, so it's easy to compare between different branches or build configs.

[dschmitz89]

Ok, so you also do not have the silver bullet ;). One env per array library sounds sensible. Not a great developer experience but unavoidable it seems. If someone with experience could add something to the developer docs, that would be great. GPU stuff could be left out in the beginning as that is a minefield we wanna avoid for some time.

[lucascolley]

I'm happy to tackle this properly in the summer (in terms of separate environments), but we should probably include instructions for venvs as well. Feel free for anyone else to write something before then. Maybe we'll even get the GPU drivers figured out (I may play around with wiping my machine / starting new containers from scratch like Andrew has talked about before).

The GPU stuff would probably be useful more widely for people who want to use a combination of CuPy/PyTorch/JAX, and I wouldn't be surprised if there are already efforts from people to document how to do it with a lot more knowledge than me.

[ilayn]

Also, have some mercy for the windows devs please. Most of the tools, in particular Jax has very little care for native environments. And I would really prefer that the array API tests are optional with a dev.py flag.

[andyfaff]

Most of the tools, in particular Jax has very little care for native environments.

Interested to know more in that regard. E.g. Are there things that escape an environment?

[rgommers]

And I would really prefer that the array API tests are optional with a dev.py flag.

They are, there's an --array-api-backend flag to dev.py test, and it requires setting an environment variable SCIPY_ARRAY_API to even enable any of this. It's also a single separate CI job only at this point.

[ilayn]

Interested to know more in that regard. E.g. Are there things that escape an environment?

Nothing particular, Jax on windows is experimental and only available for cpu. Still a bit at a loss why certain things are not running on my laptop but probably it's on me.

So if we decide to go a bit more ambitious with testing on gpu locally for devs then windows is out.

[rgommers]

We should probably also declare scipy.io as out of scope, same as scipy.datasets (xref #20594 (comment)). Given that the supported file formats are all living on disk and the I/O is inherently going through host memory, there's not really anything to do. Changing to a non-numpy array type can just as easily be done by the user with a function call like xp.asarray as with a keyword to scipy.io functions.

Furthermore, I think it's time to finalize the decision on scipy.fftpack. It's legacy and we want everyone to use scipy.fft instead - so let's not touch fftpack at all.

[lucascolley]

Furthermore, I think it's time to finalize the decision on scipy.fftpack. It's legacy and we want everyone to use scipy.fft instead - so let's not touch fftpack at all.

Agreed, I recently added "remove fftpack??" to the SciPy 2.0 wiki. IIUC the plan with legacy things is to break them out into their own repo and archive them?

[rgommers]

No such plan I think - legacy in particular means "not deprecated, we're keeping it but please don't use it for new code". Of course, once something has been in legacy for a long time and usage fades away, we are always free to reconsider and actually deprecate it. That'd be a new discussion though.