Python 中的 Numpy 数组之外

执行摘要

近年来，Python 的数组计算生态系统有机地发展，以支持 GPU、稀疏和分布式数组。这非常棒，也是去中心化开源开发中可能发生的增长的一个很好的例子。

然而，为了巩固这一增长并将其应用于整个生态系统，我们现在需要进行一些集中规划，从包需要相互了解的成对模型转向包可以通过开发和遵守社区标准协议进行协商的生态系统模型。

通过适度的努力，我们可以定义一个 Numpy API 的子集，它能很好地在所有这些地方工作，从而使生态系统能够在硬件之间更平滑地过渡。这篇文章描述了实现这一目标的机遇和挑战。

我们首先讨论两类库

实现 Numpy API 的库
使用 Numpy API 并在其之上构建新功能的库

实现 Numpy API 的库

Numpy 数组是数值 Python 生态系统的基础之一，并作为其他语言中类似库的标准模型。今天，它被用于分析卫星和生物医学图像、金融模型、基因组、海洋和大气、超级计算机模拟以及来自数千个其他领域的数据。

然而，Numpy 是在几年前设计的，其实现对于一些现代硬件已不再最优，特别是多核工作站、多核 GPU 和分布式集群。

幸运的是，其他库在这些其他架构上实现了 Numpy 数组 API

CuPy：在支持 CUDA 的 GPU 上实现 Numpy API
Sparse：为大部分为零的稀疏数组实现 Numpy API
Dask array：为多核工作站或分布式集群并行实现 Numpy API

因此，即使 Numpy 的实现不再理想，Numpy API 在后续项目中依然存在。

注意：绝大多数情况下，Numpy 实现仍然是理想的。密集内存数组仍然是常见情况。这篇博客文章讨论的是 Numpy 不理想的少数情况。

所以今天我们可以在所有 Numpy、GPU、稀疏和并行数组之间编写类似的代码

import numpy as np
x = np.random.random(...)  # Runs on a single CPU
y = x.T.dot(np.log(x) + 1)
z = y - y.mean(axis=0)
print(z[:5])

import cupy as cp
x = cp.random.random(...)  # Runs on a GPU
y = x.T.dot(cp.log(x) + 1)
z = y - y.mean(axis=0)
print(z[:5].get())

import dask.array as da
x = da.random.random(...)  # Runs on many CPUs
y = x.T.dot(da.log(x) + 1)
z = y - y.mean(axis=0)
print(z[:5].compute())

...

此外，每个深度学习框架（TensorFlow、PyTorch、MXNet）都有一个类似 Numpy 的东西，它与 Numpy 的 API 有点相似，但肯定不是力求完全匹配。

使用和扩展 Numpy API 的库

在开发适用于不同硬件的 Numpy API 的同时，许多库今天在 Numpy API 的基础上构建算法功能

XArray：用于带标签和索引的数组集合
Autograd 和 Tangent：用于自动微分
TensorLy：用于高阶数组分解
Dask array：协调多个类似 Numpy 的数组形成一个逻辑并行数组

(Dask array 既使用也实现 Numpy API)
Opt Einsum：用于更高效的爱因斯坦求和运算
...

这些项目以及更多项目增强了 Python 中的数组计算能力，构建了超越 Numpy 本身提供的新功能。

还有一些项目，如 Pandas、Scikit-Learn 和 SciPy，它们使用了 Numpy 的内存内部表示。在这篇博客文章中，我们将忽略这些库，而重点关注那些只使用高级 Numpy API 而不使用低级表示的库。

机遇与挑战

鉴于这两类项目

实现 Numpy API 的新库 (CuPy, Sparse, Dask array)
使用和扩展 Numpy API 的新库 (XArray, Autograd/tangent, TensorLy, Einsum)

我们希望将它们结合使用，例如将 Autograd 应用于 CuPy，将 TensorLy 应用于 Sparse 等等，包括所有未来可能出现的实现。这具有挑战性。

不幸的是，尽管所有数组实现的 API 都与 Numpy 的 API 非常相似，但它们使用了不同的函数。

>>> numpy.sin is cupy.sin
False

这给使用方库带来了问题，因为它们现在需要根据接收到的类似数组的对象来切换使用的函数。

def f(x):
    if isinstance(x, numpy.ndarray):
        return np.sin(x)
    elif isinstance(x, cupy.ndarray):
        return cupy.sin(x)
    elif ...

今天，每个数组项目都实现了自定义插件系统，用于在一些数组选项之间进行切换。如果您有兴趣，以下是这些插件机制的链接

例如，XArray 可以使用 Numpy 数组或 Dask 数组。这对该项目的用户非常有益，他们今天可以在笔记本电脑上的小型内存数据集与集群上的 100TB 数据集之间无缝切换，全部使用相同的编程模型。然而，当考虑将稀疏或 GPU 数组添加到 XArray 的插件系统时，很快就清楚地发现今天这样做成本会很高。

构建、维护和扩展这些插件机制是昂贵的。每个项目中的插件系统都不相同，因此任何新的数组实现都必须去每个库重复构建相同的代码。同样，任何新的算法库都必须为每个 ndarray 实现构建插件。每个库都必须显式导入和理解其他库，并随着这些库随时间变化进行调整。这种覆盖是不完整的，因此用户对其应用程序在不同硬件之间的可移植性缺乏信心。

成对的插件机制对于单个项目来说是有意义的，但对于整个生态系统来说并不是一个高效的选择。

解决方案

我今天看到两种解决方案

构建一个新库，其中包含所有相关 Numpy 函数的可分派版本，并说服所有人内部使用它而不是 Numpy
将此分派机制构建到 Numpy 本身中

每种方法都有挑战。

构建一个新的集中式插件库

我们可以构建一个新的库，此处称为 arrayish，其中包含所有相关 Numpy 函数的可分派版本。然后我们说服所有人内部使用它而不是 Numpy。

因此，在每个类似数组的库的代码库中，我们编写如下代码

# inside numpy's codebase
import arrayish
import numpy
@arrayish.sin.register(numpy.ndarray, numpy.sin)
@arrayish.cos.register(numpy.ndarray, numpy.cos)
@arrayish.dot.register(numpy.ndarray, numpy.ndarray, numpy.dot)
...

# inside cupy's codebase
import arrayish
import cupy
@arrayish.sin.register(cupy.ndarray, cupy.sin)
@arrayish.cos.register(cupy.ndarray, cupy.cos)
@arrayish.dot.register(cupy.ndarray, cupy.ndarray, cupy.dot)
...

Dask、Sparse 和任何其他类似 Numpy 的库也依此类推。

在所有算法库（如 XArray、autograd、TensorLy 等）中，我们使用 arrayish 而不是 Numpy

# inside XArray's codebase
# import numpy
import arrayish as numpy

这与之前的插件解决方案相同，但现在我们构建了一个社区标准插件系统，希望所有项目都能同意使用。

这将维护多个插件系统的巨大 n 乘以 m 成本，降低到在每个库中使用单个插件系统的更易于管理的 n 加 m 成本。这个集中式项目可能也会受益于比任何单个项目自行维护得更好。

然而，这也有成本

让许多不同的项目同意一个新的标准是很困难的
算法项目将需要在内部开始使用 arrayish，添加如下新的导入
```
import arrayish as numpy
```
这肯定会在内部造成一些复杂性
需要有人构建和维护中心基础设施

Hameer Abbasi 在这里构建了一个 arrayish 的初步原型：github.com/hameerabbasi/arrayish。在 pydata/sparse #1 中，也有一些关于这个话题的讨论，使用了 XArray+Sparse 作为例子。

从 Numpy 内部进行分派

或者，中心分派机制可以位于 Numpy 内部。

Numpy 函数可以学习将其控制权交给其参数，允许数组实现尽可能接管。这将使得现有的 Numpy 代码能够在外部开发的数组实现上工作。

这有先例。 array_ufunc 协议允许任何定义了 __array_ufunc__ 方法的类接管任何 Numpy ufunc，例如 np.sin 或 np.exp。Numpy 的归约（reduction）函数，例如 np.sum，已经会查找其参数上的 .sum 方法，如果可能就委托给它们。

一些数组项目，如 Dask 和 Sparse，已经实现了 __array_ufunc__ 协议。CuPy 也有一个开放的 PR。这里有一个示例清晰地展示了 Numpy 函数在 Dask 数组上的使用。

>>> import numpy as np
>>> import dask.array as da

>>> x = da.ones(10, chunks=(5,))  # A Dask array

>>> np.sum(np.exp(x))             # Apply Numpy function to a Dask array
dask.array<sum-aggregate, shape=(), dtype=float64, chunksize=()>  # get a Dask array

我建议所有与 Numpy API 兼容的数组项目都实现 __array_ufunc__ 协议。

这适用于许多函数，但不是全部。其他操作，如 tensordot、concatenate 和 stack 在算法代码中经常出现，但此处未涵盖。

这个解决方案避免了上面 arrayish 解决方案的社区挑战。每个人都习惯于遵从 Numpy 的决定，并且相对较少的代码需要重写。

这种方法的挑战在于，历史上 Numpy 的发展速度比生态系统的其他部分慢。例如，上面提到的 __array_ufunc__ 协议在合并之前讨论了好几年。幸运的是，Numpy 最近获得了资金，以帮助其更快地进行此类更改。尽管在这笔资金下雇佣的全职开发人员刚刚开始工作，但目前尚不清楚这项工作对他们来说有多大的优先权。

就我而言，我更倾向于看到这种 Numpy 协议解决方案得到实施。

结语

社区以前也进行过这种过渡（Numeric + Numarray -> Numpy，Scikit-Learn 的 fit/predict API 等），通常取得了令人惊讶的积极成果。

我今天面临的开放性问题如下

Numpy 在满足对协议的需求的同时，如何迅速适应，并保持其作为生态系统基础的现有角色的稳定性
有哪些算法领域可以以跨硬件的方式编写，仅依赖于高级 Numpy API，而不需要在数据结构层面进行特殊化。显然存在一些这样的领域（XArray，自动微分），但这些领域有多普遍？
一旦标准协议到位，还会出现哪些其他类似数组的实现？内存压缩？概率式？符号式？

更新

在 BIDS 参加五月 NumPy 开发者冲刺活动并讨论了此话题后，我们中的几位起草了一份 Numpy 改进提案 (NEP)，可在此处获取。

博客评论由 Disqus 提供支持

Python 中的 Numpy 数组之外 为 GPU、分布式和稀疏数组准备生态系统