HackMD - Collaborative Markdown Knowledge Base

In this blog post I talk about the work that I was able to accomplish during my internship at Quansight Labs and the efforts being made towards making array libraries more interoperable. Going ahead, I'll assume basic understanding of array and tensor libraries with their usage in the Python Scientific and Data Science software stack. As someone who's been involved in the Data Science domain for more than four years now, I can't think of any of my projects without using an array/tensor library. Over time, `NumPy` (more than 15 years old today) has become a blind import before starting any Python project involving manipulation of arrays. ***\*\*\*\*\*\*\*Enter Boodha Numpy Image\*\*\*\*\*\*\*\**** Quite recently with the Deep Learning buzz, more and more libraries have been developed with some extra sweetness on top of the NumPy like nd-array computation functionality. Namely this extra sweetness is the multi-device (`CPU`, `GPU`, `XLA` etc.) support and autograd engine for the array/tensor object. One such framework is `PyTorch`, and it has been my go-to for research in machine learning during the last couple of years. The reason simply is due to this extra saccharinity along with some other machine learning specific modules (`torch.nn`, `torch.optim` etc.) it has to offer. A few eminent array/tensor libraries include CuPy, Tensorflow, MXNet, JAX etc. Henceforth I'll just call all of them **"array provider"** libraries. <img alt="Ninja Turtles Array Libraries" src="https://i.ibb.co/5jG2FBy/ninja-arrays.png"> Old Ninja Turtle NumPy and young Tensor Turtles Libraries like [`SciPy`(https://github.com/scipy/scipy), [`scikit-learn`](https://github.com/scikit-learn/scikit-learn), [`einops`](https://github.com/arogozhnikov/einops) etc. are built on top of these "array providers" and they integrate specific features through their own methods which conventionally ingest these array/tensors as their arguments. We'll refer to these set of libraries as **"array consumer"** libraries. ## What's the catch? The pain now is that an extremely stable and popular consumer library like SciPy is cut off from the sweet features which the new array/tensor frameworks have to offer. Albeit a user might prefer SciPy's domain specific features and stable API for their scientific stack, but because the underlying project is not using NumPy, and is built on top of say PyTorch, they are left at sea and are unable to use SciPy. After some thought, one may end up choosing either of the following options. 1. Convert the inputs to NumPy before feeding them in the SciPy function which is where you lose the "extra sweetness". 2. Find some other consumer library (if exists) mimicking SciPy's API but built on top of the array/tensor library that was used in the initial project (`PyTorch`, in this case). This way the user enjoys the extra features, but this comes at the cost of learning and getting used to the new library API which mimics SciPy functionality. The APIs (function signatures) for these "array providers" are also not the same all the time which is also a problem in itself. Imagine a scientiest who has experience using NumPy for about 10 years. When they move into the PyTorch ecosystem, they are still doing a lot of similar computations, but now will need to learn a new API. Life is hard, ain't it! But as Rocky Balboa said ([actual quote :P](https://youtu.be/8xFEqdkO5UI?t=13)) > "NumPy, CuPy, PyTorch or SciPy is not gonna hit as hard as all of them when used together. But it ain't about finding a way to use them individually; it's about making them work together. That's how interoperable science is done." OK sorry, that's just me using Rocky's motivational lines to make a point. To define the issue more concretely, the question is, can we do something like the following?! ```python import torch ... # Do something with PyTorch x = torch.rand(3000) # End up with some torch tensor # Estimate power spectral density using Welch’s method. # SciPy offers a method `welch` for doing exactly that. from scipy.signal import welch f, Pxx = welch(x, fs=10e3, nperseg=1024) ``` Our world would be a better place if we could pass any kind of array/tensor object in these consumer libraries and let them do their magic, finally spitting out the same array/tensor type as the input. (Can rephrase this?) What if there is a way? What if there are existing ways to achieve this? Not kidding, it is possible with recent efforts made towards this direction. ## Protocols for Interoperable Behaviour Now that the problem and motivation is pellucid, let's dive into array libraries interoperability and understanding the protocols making this a reality. Enter [NEP 18 (`__array_function__` protocol)](https://numpy.org/neps/nep-0018-array-function-protocol.html), which was one of the first efforts to address the issue of interoperability in NumPy. In a nut shell NEP 18 allows arguments of NumPy functions to define how that function operates on them. This enables using NumPy as a high level API for efficient multi-dimensional array operations, even with array implementations that differ greatly from `numpy.ndarray`. I suggest reading the [NEP](https://numpy.org/neps/nep-0018-array-function-protocol.html) itself for detailed understanding, but I'll try to expound the motivation with a simple example taken from [this](https://www.youtube.com/watch?v=HVLPJnvInzM) insightful talk by [Dr. Ralf Gommers](https://github.com/rgommers) at PyData Amsterdam. ```python import numpy as np import cupy def some_func(x): return np.log(np.sin(x) + 5) x = np.random.random((100000, 100000)) some_func(x) # Runs on CPU, might be slow # Now can we use some_func with CuPy arrays # designed to work on NVIDIA GPUs and are # faster at parallelized matrix operations. x_cupy = cupy.array(x) # NEP 18 enables this possibility some_func(x_cupy) # Runs on GPU, orders of magnitude fast ``` Since NEP 18, there have been a few other protocols like [NEP 35](https://numpy.org/neps/nep-0035-array-creation-dispatch-with-array-function.html), [NEP 37](https://numpy.org/neps/nep-0037-array-module.html) endeavouring to address some of the issues and shortcomings with NEP 18. For the sake of this blog, we'll limit our focus on [NEP 31 or `uarray`](https://uarray.org/en/latest/) and [NEP 47 or Array API (`__array_namespace__`)](https://numpy.org/neps/nep-0047-array-api-standard.html) which are some of the most recent protocols with a goal to ameliorate interoperability. ## Array API or NEP 47 Before I start describing Python Array API in the context of this blog, I urge you to read: 1. [Announcing the consortium](https://data-apis.org/blog/announcing_the_consortium/) 2. [First release of the Array API Standard](https://data-apis.org/blog/array_api_standard_release/) These two official blogs from the [Consortium for Python Data API Standards](https://data-apis.org/) describe the protocol, giving a high level overview of its existence. With the introduction and compliance of Array API in all the major array/tensor libraries, consumer libraries will be able to support more than one array framework and become truly interoperable. Let's see how the use of Array API might look in practice. ```python import torch as xp # or import numpy.array_api as xp a = xp.arange(3) b = xp.ones(3) c = a + b ``` Technically, the only changes involved for a NumPy end user to port their code in PyTorch would be to change their import statement and refactor: `np.` -> `xp.`. Where `xp` represents any array/tensor implementation compliant with the Array API. Doing something like this is duck soup compared to some other array interoperability protocols taking a much more convoluted approach. A couple of the concrete use cases mentioned in the [Array API Spec](https://data-apis.org/array-api/latest/use_cases.html#concrete-use-cases) are: * Use case 1: add hardware accelerator and distributed support to SciPy * Use case 2: simplify einops by removing the backend system Since we started by talking about consumer libraries like SciPy, let's continue with the same example. We've built a [demo](https://anirudhdagar.ml/array-api-demo/intro.html) around Array API showcasing the use of PyTorch Tensors with Scipy. The [Array API Demo](https://anirudhdagar.ml/array-api-demo/intro.html) walks you through the details and processes involved to make an array consumer library like SciPy more interoperable with array provider libraries. The demo is built keeping two different perspectives in mind, one that of an end user and the other of an open-source developer/maintainer looking to incorporate and use Array API within their array consumer library. In essence, the demo is built on top of the 2017 Nobel prize winning work for Physics about [LIGO-Virgo detector noise and extraction of transient gravitational-wave signals](https://inspirehep.net/literature/1751757). The original tutorial ([collab link](https://colab.research.google.com/github/losc-tutorial/Data_Guide/blob/master/Guide_Notebook.ipynb)) was built using NumPy, SciPy, Matplotlib and other open source libraries in the PyData ecosystem. But, instead of using NumPy arrays we switch to a PyTorch based implementation with minimal changes in the codebase. Let's dive into the exact formulation and python code that allows this behaviour. ### get_namespace The demo shows the implementation of a dummy [`get_namespace`](https://anirudhdagar.ml/array-api-demo/GW_Demo_Array_API.html#get-namespace) method which is the first function to be called inside any SciPy method. One can see how it works below, simply returning the namespace, which can be used later for calling any Array API specified methods. ```python def get_namespace(*xs): # `xs` contains one or more arrays, or possibly Python scalars (accepting # those is a matter of taste, but doesn't seem unreasonable). namespaces = { x.__array_namespace__() if hasattr(x, '__array_namespace__') else None for x in xs if not isinstance(x, (bool, int, float, complex)) } if not namespaces: # one could special-case np.ndarray above or use np.asarray here if # older numpy versions need to be supported. # This can happen when lists are sent as an input, for eg. some # SciPy functions coerce lists into ndarrays internally. raise ValueError("Unrecognized array input") if len(namespaces) != 1: raise ValueError(f"Multiple namespaces for array inputs: {namespaces}") xp, = namespaces if xp is None: raise ValueError("The input is not a supported array type") return xp ``` This should be possible with the Array API, for the end-user in a future release of SciPy and PyTorch. Please read more on the [demo](https://anirudhdagar.ml/array-api-demo/intro.html) page itself. Interestingly in the process, with an effort to make PyTorch more [array api](https://data-apis.org/array-api/latest/) compatible, I also started contributing to PyTorch. My first PyTorch PR [#62560](https://github.com/pytorch/pytorch/pull/62560) started with adding an alias `torch.concat` for `torch.cat` which is how concatenation is defined in the array api spec and ended with fixing a bunch of other things related to concatenation in PyTorch. I worked on adding and improving compatibility for Array API to the `torch.linalg` module as well recently. See [pytorch/pytorch#63285](https://github.com/pytorch/pytorch/pull/63285) and [pytorch/pytorch#63227](https://github.com/pytorch/pytorch/pull/63227). ## uarray ``uarray`` is a backend system for Python that allows you to separately define an API, along with backends that contain separate implementations of that API. For uarray I've been working on adding support to more SciPy modules. This [uarray backend compatibility tracker](https://github.com/scipy/scipy/issues/14353) in SciPy issue sums up the plan and current state of uarray in SciPy. Precisely I've been working on adding uarray support in the `ndimage` module for the last couple months. See [scipy/scipy#14356](https://github.com/scipy/scipy/pull/14356) for more details on the discussion. With the `ndimage` ```python import numpy as np import cupy def some_func(x): return np.log(np.sin(x) + 5) x = np.random.random((100000, 100000)) some_func(x) # Runs on CPU, might be slow # Now can we use some_func with CuPy arrays # designed to work on NVIDIA GPUs and are # faster at parallelized matrix operations. x_cupy = cupy.array(x) # NEP 18 enables this possibility some_func(x_cupy) # Runs on GPU, orders of magnitude fast ``` From the perspective of interoperability I list below some Pros & Cons for these two protocols. * `uarray` * It can handle compiled code. * Involves lot of setup machinery to make it work * The dispatching mechanism is not implicit. The user needs to register the backend before he/she can use the array library of choice. See the disucssion we opened for auto registering the backend but that is probably something maintainers don't really want. * Array API & `__array_namespace__` protocol * Extremely easy for libraries to add support to array api and make their methods compatible with multiple "array providers". * One can make use of the consumer library functions without thinking about any other registration etc. and passing the array implementation of their choice. Additionally array api also makes it easier for the user to switch between different array implementations since the signatures are standardized across frameworks. * Can't handle compiled code, works for python/numpy based consumer library methods only. ## What's Next? I've absolutely enjoyed the Scientefic Python Open-Source community and plan to continue working on projects including SciPy and PyTorch voluntarily going forward. More specifically, I plan to work on improving interoperability with other libraries in `PyTorch` with `Python Array API` compliance, which is aimed for a release in `1.11` and also improving NumPy support. There are a lot of interesting gaps that I aim to fill in the `OpInfo` testing module and in general trying to catch a few bugs through the improved testing framework. With the recent migration to `Structured Kernels`, I also plan to help out with porting of some Ops to `structured`. In `SciPy`, I aim to continue adding and improving `uarray` backend support for more modules. This also extends work into libraries like `CuPy`, `Dask` etc. where the `uarray` backend needs to be enabled. Apart from `uarray`, I'd also like to explore and contribute to a few more intersting tracks currently being focussed on in SciPy. These include the work involving ## References * [Python Data API Standards](https://data-apis.org/) * [uarray](https://uarray.org/en/latest/) * [NEP 18 — A dispatch mechanism for NumPy’s high level array functions (`__array_function__` Protocol)](https://numpy.org/neps/nep-0018-array-function-protocol.html) * [NEP 31 — Context-local and global overrides of the NumPy API](https://numpy.org/neps/nep-0031-uarray.html) * [NEP 47 — Adopting the array API standard](https://numpy.org/neps/nep-0047-array-api-standard.html) * [Ralf Gommers: The evolution of array computing in Python | PyData Amsterdam 2019](https://www.youtube.com/watch?v=HVLPJnvInzM)