Array Libraries Interoperability

Anirudh Dagar

This talk on Github: anirudhdagar/qs-intern-talk

Array Providers





JAX etc.

Array Consumers





einops etc.

NumPy, CuPy, PyTorch or SciPy aren’t 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. Rocky Balboa

Try using non-NumPy with SciPy

# Estimate power spectral density using Welch’s method.
# SciPy offers a method `welch` for doing exactly that.
from scipy.signal import welch

freq, Pxx = welch(x)

# For CuPy
ValueError: object __array__ method not producing an array

# For PyTorch
TypeError: can't convert cuda:0 device type tensor to numpy.

Python Array API

  • Standard Array API across all array/tensor frameworks.
  • NEP 47 or __array_namespace__ protocol

Python Array API: Use Case

  • add hardware accelerator and distributed support to SciPy
  • simplify einops by removing the backend system

Array API Demo: SciPy with PyTorch Tensors

PyTorch Contributions

Improve Array API Compliance in PyTorch


  • backend / dispatch mechanism
  • separately define an API; backends contain separate implementations of that API.
  • user can register an implementation

uarray Tracker : Issue #14353

SciPy uarray ndimage PR:

# SciPy ndimage with CuPy array
from scipy import ndimage
import cupy as cp

with scipy.ndimage.set_backend('cupy'):
    y_cupy = ndimage.correlate1d(cp.arange(10),
                                 cp.array([1, 2.5]))

uarray Demo: SciPy with CuPy Arrays


Anirudh Dagar

This talk on Github: anirudhdagar/qs-intern-talk