The Array API Standard in SciPy
Published October 4, 2023
In this blog post I discuss how SciPy is harnessing the power of array-agnostic code to become more interoperable, and why that matters.
Hi 👋 I'm Lucas Colley, a Computer Science & Philosophy student from the UK. This summer, I have had the amazing opportunity of working under the mentorship of Ralf Gommers, Pamphile Roy and Irwin Zaid in my Open Source internship at Quansight Labs. My project was to make progress towards SciPy's goal of adopting support for the array API standard, which will allow users to use SciPy functions with CuPy arrays, PyTorch tensors and potentially many other array/tensor libraries in the future.
This post stands on the shoulders of giants, including Anirudh Dagar's blog post "Array Libraries Interoperability" (from which my illustration takes much inspiration!) and Aaron Meurer's SciPy 2023 talk "Toward Array Interoperability in the Scientific Python Ecosystem"; for details of wider interoperability efforts and the array API standard, I urge you to check out those links!
Here, I approach the topic from the focused perspective of a SciPy contributor. I hope that this post will be helpful for anyone who wants to contribute to SciPy's efforts to support the array API standard.
SciPy and the Scientific Python Ecosystem
SciPy is one part of a wide collection of Open Source tools, written in Python, which are fundamental to a range of work which relies on manipulating numbers with computers. This collection is known as the Scientific Python Ecosystem - a term which broadly encompasses Open Source Python projects which are community driven and used for science (see scientific-python.org for more information).
NumPy is the n-dimensional array library which SciPy is built to work with - it is SciPy's only non-optional runtime dependency. At the time of writing, the vast majority of SciPy functions only work when providing NumPy arrays, despite the various interoperability efforts so far (see Anirudh's blog post). The main problem with this is that "NumPy provides an array implementation that’s in-memory, CPU-only and single-threaded." The potential performance improvements that could be won with parallel algorithms, support for distributed arrays and support for GPUs (and other hardware accelerators) are huge, and thus "SciPy has had 'support for distributed and GPU arrays' on its roadmap" for five years now. For more information, see this page on the use cases of the array API, and the blog post: "A vision for extensibility to GPU & distributed support for SciPy, scikit-learn, scikit-image and beyond" by Ivan Yashchuk and Ralf Gommers.
Due to these performance wishes, many different 'array provider' libraries, such as CuPy and
PyTorch, have been developed, with support for one or more of these hardware use-cases in mind.
However, due to the fact that 'array consumer' libraries are typically built to be compatible with
only one array provider library, many functions have been reimplemented in different consumer libraries,
just to work with a different provider library.
In extreme cases, this has led to efforts to reimplement entire consumer libraries,
This constant reimplementation of the same functionality has led many to see the Scientific Python Ecosystem
as rather being fractured into many smaller ecosystems, one built around each array library, and to dream for a
world where the ecosystems are unified and everything works together!
While it would certainly be a very ambitious goal to aim for everything working together, the goal of increasing interoperability between array libraries is one with widespread support. There are currently many hurdles which users can run into when navigating this fractured ecosystem. They may find that they want to convert between different array provider libraries in their projects, in order to make use of the full range of tools which have been developed. However, for many projects, frequent conversion between provider libraries is a complete no-go: cross-device copying (e.g. from GPU to CPU when converting from CuPy to NumPy) can be a huge performance hit, and other features, such as automatic differentiation, can be lost in conversion.
For users who decide to use a different provider library to take advantage of performance benefits, they can find themselves needing to learn new APIs to accomplish the same tasks, which are often similar but frustratingly different to those which they are used to. There are also likely to be gaps in the functionality offered by the libraries, which can lead back to the problem of converting between libraries, or require the user to find new ways to complete the same task. For large existing projects, this makes the prospect of switching to a different provider library a very daunting one.
On the developer side of things, it seems like there is a lot of wasted time in reimplementing the same functionality for many different provider libraries. Developer time is very valuable since these projects are still largely maintained by volunteers (and hence large tasks like adding GPU support can stay on the roadmap for many years), so work to reduce the overall maintenance load is key to accelerating progress. If each piece of functionality only needed to be implemented once and worked with any array library, then there would be many fewer projects to maintain, and much more time to work on implementing brand new functionality. Fortunately, there has been work on a standard which will help push us towards this imagined future.
The Array API Standard
The Consortium for Python Data API Standards has created the Python array API standard, which is a "specification for what it means for a library to have an ‘array API’" (see Aaron's talk). The Consortium was formed in 2020 by a group of maintainers and industry stakeholders - see the paper accompanying Aaron's talk for more information. Array consumer libraries can write 'array-agnostic' code which is compatible with this specification, rather than any particular provider library's API, and it will work with any provider library which complies with the standard. NumPy, CuPy and PyTorch are all planning to be fully compliant - with this being one of the goals for NumPy 2.0 - and JAX and Dask also have implementations in progress. The consortium has also developed array-api-compat, a small wrapper around each array library (currently NumPy, CuPy and PyTorch), which, for practical purposes, is fully compliant with the standard.
This means that array consumer libraries are able to start writing and testing array-agnostic code today.
For a NumPy-consuming library like SciPy, this has opened the opportunity to implement support for CuPy and
PyTorch by adding support for the standard via
SciPy support for the array API
Before these interoperability efforts started, SciPy did not work well with provider libraries other than
NumPy at all (as expressed by
this amazing illustration
from Anirudh's blog post).
For example, when trying to use
scipy.fft.fft with a CuPy array, users would be hit with this error message:
Adopting the standard is the first step towards a future where users who have programs made up of SciPy functions can simply change the provider library used for their input data, and have the rest of the program still work the same, while taking advantage of any benefits of the new provider library.
To kick off the process of this work, the request-for-comment (RFC) issue
"SciPy array types & libraries support" was opened.
This set out the basic design principle of "container type in == container type out", as well as a
detailed plan for how to treat different array inputs, and the development strategy.
To summarise, the scope of the work is to treat all array inputs as specified in the RFC and to convert all
existing pure Python + NumPy code to be array API compatible.
Out of scope is converting the C/C++/Cython/Fortran code which is used within SciPy.
Behind the scenes, SciPy provides Python APIs for low level functions which are written in other languages and for the
These APIs are made to work with NumPy arrays - they will definitely not work for arrays on a different device to the CPU,
and are not guaranteed to work with other types of array on the CPU.
For these functions which use compiled code, some dispatching mechanism seems like the right solution, but no particular
implementation or development strategy has been settled on yet - see the discussion of
Anirudh's blog post and
the RFC for more information.
Near the start of my internship,
the first PR to convert a SciPy submodule, by Pamphile Roy, one of my mentors,
was merged, with
scipy.cluster being covered (modulo a few minor follow-ups).
This PR added a lot of machinery (which I will explore in more detail later) to enable conversion and testing,
while keeping the new behaviour behind an experimental environment variable.
My main goal was to convert
scipy.fft, SciPy's Fast Fourier Transform module.
fft is an
array API standard extension module,
which means that array provider libraries supporting the standard can choose to implement it or not.
scipy.fft's backend is mostly written in C++, but thanks to the extension module, we were able to convert all
of the standard extension functions to be compatible with all libraries which:
- comply with the standard
- implement the
After the main PR merged, the ability to use CuPy arrays and PyTorch tensors
with the standard extension functions was enabled when the experimental environment variable
SCIPY_ARRAY_API is set to
For example, the code in the following benchmark - the task of smoothing an image (taking a 2-D image,
scipy.fftn, zero-ing out the highest frequency components, and using
scipy.ifftn) using CuPy
arrays - now works:
On my machine (AMD Ryzen 5 2600 & NVIDIA GeForce GTX 1060 6GB), this demonstrated a ~15x performance improvement over using NumPy arrays in the same functions. You can see the full benchmark on the PR.
Completing this PR wasn't without difficulty, as lots of SciPy's brand new array API testing infrastructure
needed improvements in the process.
I contributed quite a few PRs to help make this testing infrastructure more robust and squash some bugs;
you can find all of my merged contributions relating to array types
in this filtered view of SciPy's PR tracker.
This was also my first ever big contribution to Open Source, so there was a large (but very rewarding!) learning
process with Git, GitHub and collaboration with reviewers.
In the next section, I'll dive into the process of converting
scipy.fft and share some of the lessons learned
for converting other submodules.
How to convert a submodule
I'll start by showing what the necessary changes look like, before explaining in more detail how some of the machinery works. Production functions can be divided into three categories, each with different conversion steps: those which just use pure Python and NumPy, those which use compiled code, and those which implement functions found in standard extension modules.
Pure Python + NumPy code
For functions which are made up of pure Python and NumPy, the conversion process is really quite simple!
It involves comparing the currently used functions/methods with those in
the standard API specification
and mimicking the current behaviour.
The key step is the addition of
xp = array_namespace(x, y), where
y (and so on for more inputs)
are the input arrays.
This line is used at the start of every converted function, and allows us to use the functions/methods which
are compatible with the arrays without re-importing anything!
The rest of the changes are mostly replacing uses of
xp, alongside making adjustments where the
standard API differs from the NumPy API.
For example, here are the modifications made to
For functions which hit compiled code, the standard isn't enough anymore.
For now, we attempt to convert to NumPy using
np.asarray before the compiled code, then convert back to our
array's namespace using
xp.asarray after it.
This will raise exceptions for arrays on a different device to the CPU, as justified in
the RFC (tldr: we want to avoid silent device transfers).
Here is an example from
scipy.fft's Discrete Sine and Cosine Transforms, where
_pocketfft is the backend
which SciPy uses for NumPy arrays:
As a side note, there has also been work to support the calling of some specific backends instead of hitting
compiled code in
scipy.special (in this PR),
which enables GPU support for at least some provider libraries.
I will keep this blog post focused on just array-agnostic code since this work in
special should ultimately be
replaced by a more robust dispatching mechanism, whereas the array-agnostic code is here to stay!
Using standard extension modules
Despite this limitation on GPU support, as mentioned above,
fft is a standard extension module, so the standard can still help us with some
For the functions in the extension, we can check whether the array's namespace implements the extension with
if hasattr(xp, 'fft'), and use
xp.fft if so.
Since NumPy has an
fft submodule, we need to be careful to still use SciPy's implementation for NumPy arrays,
for which we have the
Here is how that looks for
After converting the production functions, we also need to convert the tests.
We have the
@array_api_compatible decorator, which parametrizes
xp with the
available backends that we have for testing.
We have also written three new array-agnostic assertions -
which incorporate checks for matching namespaces and dtypes.
If any tests use results or input data from NumPy, we convert everything to the
xp namespace before using
The assertions make use of functions from
xp where possible which helps to minimise the amount of device transfers.
Other than that, we make similar changes to those in the production functions so that every test works
with all of the backends with which we test.
Finally, for tests of functions which hit compiled code, we add the
Here is what that looks like for one of the Discrete Sine and Cosine Transforms tests:
Enabling and testing with non-NumPy arrays
I'll talk a bit about implementation details now - feel free to skip this section if you are just here for an overview!
When we call
array_namespace, a little more is done than just returning the namespace.
If the environment variable
SCIPY_ARRAY_API is not set,
(the compatible wrapper for
numpy) is returned, regardless of the type of the input array.
This allows us to retain all pre-array API behaviour unless the user has set the environment variable,
since the arrays are caught by our
if is_numpy(xp) code path.
If it isn't set, we raise exceptions for known-bad subclasses
(as set out in the RFC),
MaskedArrays and object arrays.
After that, we retrieve the namespace using
The current testing infrastructure allows us to test with PyTorch (CPU & GPU) and CuPy through the command line.
python dev.py test -b pytorch tests with PyTorch, and
python dev.py test -b all tests with
numpy.array_api, CuPy and PyTorch.
numpy.array_api is a strict minimal implementation of the
standard, which allows us to make sure that our array-agnostic code is really array-agnostic -
if our tests pass for
numpy.array_api, then they will pass for any fully compliant library.
Testing with PyTorch on different devices can be enabled by setting the environment variable
'cuda' for Nvidia GPUs or
'mps' for MacOS devices with the Metal framework.
For more information on the implementation details, see the SciPy "Support for Array API" documentation.
At the time of writing, SciPy has two of its submodules,
(in the sense I have described above).
signal have partial coverage in progress, and some of the compiled-code functions in
have been covered, with additional calling-out to CuPy, PyTorch and JAX.
The tracker issue
should be kept up to date with progress as more submodules are converted.
We currently support CuPy and PyTorch through
array-api-compat, with Dask and JAX on the horizon.
The great thing is that whenever a provider library complies with the standard, it will work with our
converted submodules, without SciPy even needing to know that it exists!
Looking to the future
I have opened the first PR for
which covers the functions which are part of
linalg standard extension module.
In the short term, we hope to see PRs covering the rest of the submodules and continued refinement of our
Looking further ahead for SciPy, once all libraries which we want to support comply with the standard, we can drop
array-api-compat and get the namespaces directly from the input arrays (via the
As mentioned above, a general dispatching mechanism will hopefully materialise down the line to allow full support for all functions which use compiled code. Still, even without it, a great deal of interoperability is ready to be achieved now!
Finally, I'd like to say some words about my experience as an intern at Quansight Labs and as a member of the SciPy community.
Going from having never contributed to Open Source to feeling proficient as a contributor, and having significant contributions merged towards a bleeding-edge project, has been a challenging but very special experience! Regular meetings with my mentors and the community meant that I was always able to ask questions when I was confused, and learn about the array types project and SciPy as a whole. At Quansight Labs, I was part of a cohort of interns and we also met regularly to discuss our achievements and explain where we were stuck. Sometimes, just putting your difficulties into words can help you figure things out from a different perspective!
After joining the SciPy community as part of my internship, I am keen to contribute voluntarily in the future and help to sustain science in this way. Knowing that my code will be used in a package as widely depended on as SciPy is a big motivation, and I am invested in the success of the array API standard now!
I would like to thank my mentors Ralf Gommers, Pamphile Roy and Irwin Zaid for their incredible support and guidance during my internship! I feel very lucky to have had the opportunity to learn from them in such a practical way. I would also like to thank the SciPy community, especially Tyler Reddy and Matt Haberland who have been greatly involved in the review process of many of my PRs, and Melissa Weber Mendonça who hosted many of the community meetings which I attended (as well as coordinating the internship program at Quansight Labs!).
- The array API standard
- "Python Array API Standard: Toward Array Interoperability in the Scientific Python Ecosystem" SciPy 2023 Talk
- "Python Array API Standard: Toward Array Interoperability in the Scientific Python Ecosystem" Paper
- "SciPy array types & libraries support" RFC
- "Array Libraries Interoperability" Blog Post
- Tracker Issue for array API support in SciPy