Array Libraries Interoperability

Anirudh Dagar

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

Array Providers

PyTorch

NumPy

CuPy

MXNet

JAX etc.

Array Consumers

SciPy

scikit-learn

scikit-image

AstroPy

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
  • https://data-apis.org/array-api/latest/

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

uarray

  • 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: https://github.com/scipy/scipy/pull/14356

# 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

Thanks!

Anirudh Dagar

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