Python Array API Standard: Toward Array Interoperability in the Scientific Python Ecosystem

The array API standard (https://data-apis.org/array-api/) is a common specification for Python array libraries, such as NumPy, PyTorch, CuPy, Dask, and JAX.

This standard will make it straightforward for array-consuming libraries, like scikit-learn and SciPy, to write code that uniformly supports all of these libraries. This will allow, for instance, running the same code on the CPU and GPU.

This talk will covers the scope of the array API standard, supporting tooling which includes a library-independent test suite and compatibility layer, what work has been completed so far, and the plans going forward.

Transcript

1. A a ron Meurer, Qu a nsight July 14, 2023 11:25–11:55, Amphithe a ter 204 SciPy 2023, Austin, TX Python Array API Standard Tow a rd Arr a y Interoper a bility in the Scienti f ic Python Ecosystem These slides 
 https://github.com/data-apis/scipy-2023- presentation/blob/main/presentation/Slides.pdf
4. The next tools… The next array library…
5. The next array library…
6. >>> import scipy.signal >>> import cupy >>> x = cupy.asarray(...) >>> scipy.signal.welch(x) Traceback (most recent call last): ... TypeError: Implicit conversion to a NumPy array is not allowed. Please use `.get()` to construct a NumPy array explicitly. Example: =
7. The next array library…
8. The problem • SciPy (and scikit-learn, scikit-image, statsmodels, …) are written against NumPy. • They use import numpy as np. • Their code implicitly assumes NumPy functions and semantics.
9. Potential Solutions
10. Potential Solutions “Merge everything into NumPy”
11. Potential Solutions ✘Each library has its own set of strengths and APIs (targeted hardware, performance, distributed computation, etc.). ✘This is way too much scope for NumPy. ✘Sometimes their design decisions contradict one another. “Merge everything into NumPy”
12. Potential Solutions “Make every library match the NumPy API” (CuPy approach)
13. Potential Solutions ✘NumPy was designed around eager CPU computation. ✘Many NumPy semantics are inappropriate for other use- cases (GPU, distributed, lazy evaluation, compilation, …). ✘Some NumPy design decisions are simply bad and shouldn’t be replicated. “Make every library match the NumPy API” (CuPy approach)
14. Potential Solutions “Use NumPy __array_function__ dispatching”
15. Potential Solutions ✘Not everything is implemented in NumPy (e.g., specialized deep learning functions). ✘There’s more to an API than just function signatures (indexing, operators, methods, broadcasting, type promotion, …). ✘Users expect np.func(x) to return a NumPy array. “Use NumPy __array_function__ dispatching”
16. Better Solution: API Standardization “Create a speci f ication for what it means for a library to have an ‘array API’”
17. Better Solution: API Standardization “Create a speci f ication for what it means for a library to have an ‘array API’” ✓Won’t require any new runtime dependencies. ✓Won’t require array libraries to know about each other. ✓NumPy is not special. It's just another array library.
18. Better Solution: API Standardization “Create a speci f ication for what it means for a library to have an ‘array API’” ✓Can be be selective and only include APIs that are already commonly implemented. ✓Can avoid / make optional APIs that break important use-cases (compilation, multi-device, lazy evaluation, …).
19. Data APIs Consortium • Founded in May 2020. • Stakeholders from various array and dataframe libraries, industry, array-consuming libraries, and end users met regularly to discuss a standardized array API. d a t a - a pis.org
20. Array API Standard https://d a t a - a pis.org/ a rr a y- a pi/ • De fi nes a set of functions and semantics that any standards complaint array library should implement. • No dependencies or reliance on any speci fi c array library. It’s strictly a speci fi cation. • Includes ~200 functions and array methods. • Most functions are based on already- existing APIs from common array libraries like NumPy, CuPy, PyTorch, JAX, etc. Array API Standard 
 https://data-apis.org/array-api/
21. Example: mean() • Lists input/output parameters. • Speci fi es required input dtypes, output dtype, and output shape. • Support for integer dtypes (in this example) is left unspeci fi ed. • Exact precision of the output is left unspeci fi ed.
22. Array API Design Principles • Only speci fi es a minimal set of APIs/behaviors. Libraries can implement more if they want to. • Focus is on already-existing APIs. • APIs should support JIT and AOT compilation, distributed, accelerator, and lazy evaluation use-cases. • e.g., no polymorphic return types, mutation is optional, … • Consistent/clean API • e.g., uses functions instead of methods (mean(a) instead of a.mean()), functions do not accept “array-likes”, … • Implementation details like underlying storage, precision, and error handling are left unspeci fi ed.
23. Additional Semantics • Broadcasting • Indexing • Type Promotion • In-Memory Interchange Protocol (DLPack) • Device Support
24. Type Promotion • Basic set of common numeric dtypes (including integer, real fl oating-point, and complex). • Cross-kind type promotion is not required (e.g., int8 + float32). • Type promotion should not depend on array values or array shape. bool bool (Python) float (Python) float32 float64 int (Python) uint8 uint16 uint32 uint64 int8 int16 int32 int64 complex (Python) complex64 complex128
25. In-memory Interchange Protocol: DLPack • DLPack • Header-only dependency • Stable C ABI • Speci fi es zero-copy semantics • Support for multiple devices • Already implemented in most popular Python array libraries (NumPy, CuPy, PyTorch, Tensor fl ow, …)
26. Device Support • Three methods are implemented. • .device attribute on array objects. • device= keyword to creation functions. • .to_device() method on array objects. • All computations should occur on the same device as the input array. • Implicit transfers are not allowed (mixing devices should raise an exception).
27. Optional Extensions • There are two optional extension namespaces • linalg • fft • Not required as they may be di ffi cult for some implementations.
28. Status
29. Speci f ication Status • Two released versions v2021.12 and v2022.12. • v2021.12 was the initial release. • v2022.12 added complex numbers, fast Fourier transforms extension, and a few additional functions. • Current focus in 2023 is on adoption.
30. Array Library Implementation Status • Main numpy namespace is ~80% compliant. Planned to be fully compliant for NumPy 2.0. • CuPy closely follows NumPy. • PyTorch is also ~80% compliant. Plan is also for full compliance. • JAX implementation in progress (https://github.com/google/jax/pull/16099) • Dask Array implementation in progress (https://github.com/dask/dask/pull/8750) • NumPy contains a full minimal reference implementation, numpy.array_api (more on that later). • For practical purposes, array-api-compat (more on this later) is 100% compliant. Meaning you can use the array API today!
31. Adoption Status • SciPy and scikit-learn are both actively moving to adopt the array API. • scikit-learn 1.3 has experimental array API support in LinearDiscriminantAnalysis. • Other libraries are adding support: einops, … • If you want to chat about adopting the array API, come to my sprint this weekend!
32. How hard is it to convert code that uses NumPy into code that supports the array API?
33. Adoption Example: scikit-learn
34. Xc = [] for idx, group in enumerate(self.classes_): - Xg = X[y == group, :] - Xc.append(Xg - self.means_[idx]) + Xg = X[y == group] + Xc.append(Xg - self.means_[idx, :]) - self.xbar_ = np.dot(self.priors_, self.means_) + self.xbar_ = self.priors_ @ self.means_ - Xc = np.concatenate(Xc, axis=0) + Xc = xp.concat(Xc, axis=0) - std = Xc.std(axis=0) + std = xp.std(Xc, axis=0) std[std == 0] = 1.0 - fac = 1.0 / (n_samples - n_classes) + fac = xp.asarray(1.0 / (n_samples - n_classes)) - X = np.sqrt(fac) * (Xc / std) + X = xp.sqrt(fac) * (Xc / std) U, S, Vt = svd(X, full_matrices=False) - rank = np.sum(S > self.tol) + rank = xp.sum(xp.astype(S > self.tol, xp.int32)) Adopted from https://github.com/scikit-learn/scikit-learn/pull/22554 The standard only speci f ies boolean indices as the sole index and multidimensional indexing when all axes are indexed. dot() is not included in the standard. The matmul operator @ is used instead. The standard function is called concat(). The standard uses functions on the namespace rather than array methods (xp.std(a) instead of a.std()). Functions do not implicitly support non-array inputs. Here asarray() is used to convert the input to sqrt() to an array. Numerical functions only support numerical dtypes. sum() on a boolean array requires an explicit conversion. np is replaced with xp everywhere.
35. Answer: in most cases, not too hard • The biggest change is 
 
 import numpy as np xp = array_namespace(x) 
 • array_namespace(x) returns the corresponding array API compatible namespace for an array object x (more on this later). • A few minor cleanups from NumPy for portability. • The standard mostly looks like NumPy. • The minimal array API implementation (more on that later) helps for testing.
36. The end result: major speedups vs. NumPy! Benchmarks were run on Intel i9-9900K and NVIDIA RTX 2080. See our proceedings paper for more details.
37. Adoption Example: scikit-learn (from Thom a s F a n’s tools plen a ry slides Wednesd a y Morning) from sklearn import set_config from sklearn.discriminant_analysis import LinearDiscriminantAnalysis import torch set_config(array_api_dispatch=True) X_torch, y_torch = torch.asarray(...), torch.asarray(...) lda = LinearDiscriminantAnalysis() # X_trans is a PyTorch Tensor X_trans = lda.fit_transform(X_torch, y_torch)
38. Tooling
39. Tooling We developed several tools to help aid array API adoption: ๏ Compatibility Layer: array-api-compat ๏ Minimal Implementation: numpy.array_api ๏ Array API Test Suite: array-api-tests
40. Compatibility Layer: array-api-compat • array-api-compat is a small wrapper around each library to cleanup the small di ff erences in existing libraries. • Example: array_api_compat.numpy.concat() wraps numpy.concatenate(). • Small, venderable, pure Python library with no hard dependencies. • Currently wraps NumPy, CuPy, and PyTorch. • Dask array and JAX are coming soon. https://github.com/d a t a - a pis/ a rr a y- a pi-comp a t https://github.com/data-apis/ array-api-compat
41. • Example array-api-compat usage: Compatibility Layer: array-api-compat https://github.com/d a t a - a pis/ a rr a y- a pi-comp a t array_namespace(x) returns xp, the corresponding array API compatible namespace for an array object x. from array_api_compat import array_namespace def some_function(x, y): xp = array_namespace(x, y) # Now use xp as the array library namespace return xp.mean(x, axis=0) + 2*xp.std(y, axis=0) https://github.com/data-apis/ array-api-compat
42. • numpy.array_api is a strict minimal implementation of the standard • Fails on behavior not explicitly required by the standard, even if the standard allows it. • Example: Minimal Array API Implementation: numpy.array_api >>> import numpy.array_api as xp <stdin>:1: UserWarning: The numpy.array_api submodule is still experimental. See NEP 47. >>> a = xp.ones((3,), dtype=xp.float64) >>> b = xp.ones((3,), dtype=xp.int64) >>> a + b Traceback (most recent call last): ... TypeError: float64 and int64 cannot be type promoted together
43. Minimal Array API Implementation: numpy.array_api • numpy.array_api is not designed for use by end-users. • Rather, it is for array consuming libraries like SciPy or scikit-learn to test that their use of the array API is portable against the standard. • Most aspects of the standard are not strict, e.g., cross-kind type promotion is allowed but not required. numpy.array_api explicitly disallows these. • If your code runs against numpy.array_api, it will run against any array API compliant library. • Full list of di ff erences from the main numpy namespace can be found at 
 https://numpy.org/doc/stable/reference/array_api.html
44. Array API Test Suite: array-api-tests • The standard is quite large (~200 functions, each with a set of speci fi ed semantics) • Need a way for array libraries to test their level of compliance. • We developed a test suite for the standard, array-api-tests • Over 1000 tests for every function and aspect of the standard. • Has been successfully used to implement compliance in NumPy, CuPy, and PyTorch. • This is the fi rst instance (we know of) in the Python world of a test suite that is library independent. https://github.com/d a t a - a pis/ a rr a y- a pi-tests
45. Array API Test Suite: array-api-tests • Makes heavy use of the Hypothesis property-based testing library • We implemented native Hypothesis support for the array API upstream (hypothesis.extra.array_api). • Hypothesis is a great fi t for such a test suite: • Property-based tests allow roughly one-to-one translation of the standard to tests. • Automatically tests corner cases & all possible combinations. • Example: NumPy does not follow type promotion rules, but only for 0-D array + non-0-D array (see NEP 50). Hypothesis (hypothesis.works)
46. Methodology
47. Methodology • We only wanted to standardize APIs that are: 1. Already implemented in most array libraries. 2. Used heavily in the ecosystem. • We used a data-driven approach.
48. What is already implemented in array libraries? • We compared APIs across common array libraries (including NumPy, Dask Array, CuPy, MXNet, JAX, TensorFlow, and PyTorch). • Took the common subset intersection of implemented functions and their keyword arguments. • Source code with data is available at 
 https://github.com/data-apis/array-api-comparison
49. mean(a, axis=None, keepdims=False) numpy.mean(a, axis=None, dtype=None, out=None, keepdims=<no value>) cupy.mean(a, axis=None, dtype=None, out=None, keepdims=False) dask.array.mean(a, axis=None, dtype=None, out=None, keepdims=False, split_every=None) jax.numpy.mean(a, axis=None, dtype=None, out=None, keepdims=False) mxnet.np.mean(a, axis=None, dtype=None, out=None, keepdims=False) tf.math.reduce_mean(input_tensor, axis=None, keepdims=False, name=None) torch.mean(input, dim, keepdim=False, out=None) Common Signature Subset
50. How often are different APIs actually used? • We runtime instrumented the test suites of common array-consuming libraries (SciPy, pandas, Matplotlib, Xarray, scikit-learn, statsmodels, scikit-image, etc.). • Every NumPy API call was recorded, along with its input type signature. • Each API was then ranked with the Dowdall positional voting system (a variant of the Borda count that favors APIs having high relative usage). • Source code with data is available at 
 https://github.com/data-apis/python-record-api
51. @overload def mean(a: numpy.ndarray): """ usage.dask: 21 usage.matplotlib: 7 usage.scipy: 26 usage.skimage: 36 usage.sklearn: 130 usage.statsmodels: 45 usage.xarray: 1 """ @overload def mean(a: List[float]): """ usage.networkx: 6 usage.sklearn: 3 usage.statsmodels: 9 """ How often are different APIs actually used? • Example output of python-record-api for mean(). • This example shows that np.mean(x) with x as an np.ndarray is much more common than with x as a list of floats. • This sort of usage data helped us to understand, for instance, that it’s OK to omit “array-like” inputs from the standard.
52. Future Work • Our focus in 2023 is on adoption. • NumPy 2.0 (later this year) is planned to have full array API compliance. • The standard will continue to evolve to match community needs. Your feedback on what you’d like to see is welcome https://github.com/data-apis/array-api/issues/
53. • The Data APIs Consortium has been working on a similar e ff ort for dataframes. • Dataframe Interchange Protocol https://data-apis.org/dataframe-protocol/ • Work-in-progress Dataframe API standard https://data-apis.org/dataframe-api/draft/ • See the talk by my Quansight colleague Marco Gorelli at EuroSciPy later this year! https://pretalx.com/euroscipy-2023/talk/N8YGEW/ Future Work Marco Gorelli
54. Acknowledgements Consortium Members a nd People who Helped with the Arr a y API St a nd a rd https://github.com/data-apis/array-api#contributors-
55. Acknowledgements My Qu a nsight Colle a gues Ralf Gommers Quansight Labs Co-Director Matthew Barber Athan Reines Stephannie Jimenez Gacha Thomas J. Fan scikit-learn core developer Pamphile Roy SciPy core developer
56. Acknowledgements D a t a APIs Sponsors https://data-apis.org/
57. Feedback Welcome • Feedback on the Consortium work is always welcome (new ideas for standardization, requests for help with adoption, general questions, etc.) 
 https://github.com/data-apis/array-api/issues/ • Contributions are also welcome, especially helping various libraries with adoption. • I will be sprinting on the array API this weekend. If you want help 
 adopting the array API or just have questions, please come! • See our SciPy 2023 Proceedings paper for more details. Array API Feedback 
 https://github.com/data-apis/ array-api/issues/
58. Questions These slides 
 https://github.com/data-apis/ scipy-2023-presentation/blob/ main/presentation/Slides.pdf Array API Feedback 
 https://github.com/data-apis/ array-api/issues/ Array API Standard 
 https://data-apis.org/array-api/