与 NumPy 的互操作性#

NumPy 的 ndarray 对象提供了一个用于对数组结构化数据进行操作的高级 API，以及一个基于内存中步长存储的 API 具体实现。虽然此 API 强大且相当通用，但它的具体实现存在局限性。随着数据集的增长以及 NumPy 在各种新环境和架构中的使用，在某些情况下，内存中步长存储策略不合适，这导致不同的库为其自身用途重新实现了此 API。这包括 GPU 数组 (CuPy)、稀疏数组 (scipy.sparse、PyData/Sparse) 和并行数组 (Dask 数组) 以及深度学习框架中各种类似 NumPy 的实现，如 TensorFlow 和 PyTorch。类似地，许多项目构建在 NumPy API 之上，用于标签和索引数组 (XArray)、自动微分 (JAX)、掩码数组 (numpy.ma)、物理单位 (astropy.units、pint、unyt)，以及其他在 NumPy API 之上添加额外功能的项目。

然而，用户仍然希望使用熟悉的 NumPy API 来处理这些数组，并以最少的（理想情况下为零）移植开销来重用现有代码。为了实现这个目标，定义了各种协议，用于实现与 NumPy 高级 API 相匹配的多维数组。

从广义上讲，与 NumPy 的互操作性使用了三组功能

将外来对象转换为 ndarray 的方法；
将 NumPy 函数的执行延迟到另一个数组库的方法；
使用 NumPy 函数并返回外来对象实例的方法。

我们在下面描述这些功能。

1. 在 NumPy 中使用任意对象#

NumPy API 中的第一组互操作性功能允许在可能的情况下将外来对象视为 NumPy 数组。当 NumPy 函数遇到外来对象时，它们将按顺序尝试以下方法

缓冲区协议，如 Python C-API 文档中所述。
__array_interface__ 协议，如本页中所述。它是 Python 缓冲区协议的前身，它定义了一种从其他 C 扩展访问 NumPy 数组内容的方法。
__array__() 方法，它要求任意对象将其自身转换为数组。

对于缓冲区协议和 __array_interface__ 协议，对象都会描述其内存布局，而 NumPy 则负责其他所有事情（如果可能，则为零拷贝）。如果不可能，则对象本身负责从 __array__() 返回一个 ndarray。

DLPack 是一种协议，以与语言和设备无关的方式将外来对象转换为 NumPy 数组。NumPy 不会使用 DLPack 将对象隐式转换为 ndarray。它提供了 numpy.from_dlpack 函数，该函数接受实现 __dlpack__ 方法的任何对象，并输出 NumPy ndarray（通常是输入对象数据缓冲区的视图）。DLPack 的 Python 规范页面详细解释了 __dlpack__ 协议。

数组接口协议#

数组接口协议定义了一种方法，允许类似数组的对象重用彼此的数据缓冲区。其实现依赖于以下属性或方法的存在

__array_interface__：一个 Python 字典，包含类似数组对象的形状、元素类型，以及可选的数据缓冲区地址和步长；
__array__()：一个返回 NumPy ndarray 拷贝或类似数组对象视图的方法；

可以直接检查 __array_interface__ 属性

>>> import numpy as np
>>> x = np.array([1, 2, 5.0, 8])
>>> x.__array_interface__
{'data': (94708397920832, False), 'strides': None, 'descr': [('', '<f8')], 'typestr': '<f8', 'shape': (4,), 'version': 3}

__array_interface__ 属性也可以用于就地操作对象数据

>>> class wrapper():
...     pass
...
>>> arr = np.array([1, 2, 3, 4])
>>> buf = arr.__array_interface__
>>> buf
{'data': (140497590272032, False), 'strides': None, 'descr': [('', '<i8')], 'typestr': '<i8', 'shape': (4,), 'version': 3}
>>> buf['shape'] = (2, 2)
>>> w = wrapper()
>>> w.__array_interface__ = buf
>>> new_arr = np.array(w, copy=False)
>>> new_arr
array([[1, 2],
       [3, 4]])

我们可以检查 arr 和 new_arr 是否共享相同的数据缓冲区

>>> new_arr[0, 0] = 1000
>>> new_arr
array([[1000,    2],
       [   3,    4]])
>>> arr
array([1000, 2, 3, 4])

`array()` 方法#

__array__() 方法确保任何类似 NumPy 的对象（数组、任何公开数组接口的对象、其 __array__() 方法返回数组或任何嵌套序列的对象）都可以作为 NumPy 数组使用。如果可能，这意味着使用 __array__() 创建类似数组对象的 NumPy ndarray 视图。否则，这会将数据复制到新的 ndarray 对象中。这不是最佳的，因为将数组强制转换为 ndarray 可能会导致性能问题或创建复制和元数据丢失的需要，因为原始对象及其可能具有的任何属性/行为都将丢失。

该方法的签名应该是 __array__(self, dtype=None, copy=None)。如果传递的 dtype 不是 None 且与对象的类型不同，则应将类型转换为指定的类型。如果 copy 为 None，则仅在 dtype 参数强制执行时才应进行复制。对于 copy=True，应始终进行复制，而 copy=False 应该在需要复制时引发异常。

如果某个类实现了旧的签名 __array__(self)，则对于 np.array(a)，将发出警告，说明缺少 dtype 和 copy 参数。

要查看包括 __array__() 用法的自定义数组实现示例，请参阅编写自定义数组容器。

DLPack 协议#

DLPack 协议定义了内存中步长 n 维数组对象的布局。它为数据交换提供了以下语法

一个 numpy.from_dlpack 函数，该函数接受具有 __dlpack__ 方法的（数组）对象，并使用该方法构造一个包含来自 x 的数据的数组。
__dlpack__(self, stream=None) 和数组对象上的 __dlpack_device__ 方法，它们将在 from_dlpack 内部调用，以查询数组所在的设备（可能需要传递正确的流，例如在多个 GPU 的情况下）以及访问数据。

与缓冲区协议不同，DLPack 允许交换包含除 CPU 之外设备（例如 Vulkan 或 GPU）上的数据的数组。由于 NumPy 仅支持 CPU，因此它只能转换数据存在于 CPU 上的对象。但是，其他库，如 PyTorch 和 CuPy，可以使用此协议交换 GPU 上的数据。

2. 在不转换的情况下操作外来对象#

NumPy API 定义的第二组方法允许我们在不转换的情况下将 NumPy 函数的执行延迟到另一个数组库。

考虑以下函数。

>>> import numpy as np
>>> def f(x):
...     return np.mean(np.exp(x))

请注意，np.exp 是一个 ufunc，这意味着它以逐元素的方式对 ndarray 进行操作。另一方面，np.mean 沿着数组的一个轴进行操作。

我们可以直接将 f 应用于 NumPy ndarray 对象

>>> x = np.array([1, 2, 3, 4])
>>> f(x)
21.1977562209304

我们希望此函数对任何类似 NumPy 的数组对象都能同样有效。

NumPy 允许一个类通过以下接口来表明它希望以自定义的方式处理计算

__array_ufunc__：允许第三方对象支持和覆盖 ufunc。
__array_function__：对于 __array_ufunc__ 协议没有涵盖的通用函数，这是一个万能方法。

只要外来对象实现了 __array_ufunc__ 或 __array_function__ 协议，就可以在不需要显式转换的情况下对它们进行操作。

`__array_ufunc__` 协议#

通用函数（或简称为 ufunc）是一个“矢量化”的包装器，用于函数，该函数接受固定数量的特定输入并产生固定数量的特定输出。如果并非所有输入参数都是 ndarray，则 ufunc（及其方法）的输出不一定是 ndarray。实际上，如果任何输入定义了 __array_ufunc__ 方法，则控制权将完全传递给该函数，即 ufunc 将被覆盖。该（非 ndarray）对象上定义的 __array_ufunc__ 方法可以访问 NumPy ufunc。由于 ufunc 具有明确定义的结构，因此外来 __array_ufunc__ 方法可以依赖于 ufunc 属性，如 .at()、.reduce() 等。

子类可以通过覆盖默认的 ndarray.__array_ufunc__ 方法来覆盖在它上面执行 NumPy ufunc 时发生的事情。此方法将执行，而不是 ufunc，并且应该返回操作的结果，或者如果请求的操作未实现，则返回 NotImplemented。

`__array_function__` 协议#

为了实现对 NumPy API 的足够覆盖以支持下游项目，需要超越 __array_ufunc__，并实现一个协议，该协议允许 NumPy 函数的参数控制并以一种安全且跨项目一致的方式将执行转移到另一个函数（例如，GPU 或并行实现）。

__array_function__ 的语义与 __array_ufunc__ 非常相似，只是操作是通过一个任意的可调用对象而不是一个 ufunc 实例和方法来指定的。有关更多详细信息，请参阅 NEP 18 — NumPy 高级数组函数的调度机制。

3. 返回外来对象#

第三种类型的功能集旨在使用 NumPy 函数实现，然后将返回值转换回外来对象的实例。 __array_finalize__ 和 __array_wrap__ 方法在幕后起作用，以确保 NumPy 函数的返回类型可以根据需要指定。

__array_finalize__ 方法是 NumPy 提供的一种机制，允许子类处理创建新实例的各种方法。每当系统从 ndarray 的子类（子类型）的对象内部分配一个新数组时，都会调用此方法。它可以用于在构造后更改属性，或从“父级”更新元信息。

The __array_wrap__ 方法在某种意义上“封装了操作”，允许任何对象（如用户定义函数）设置其返回值的类型并更新属性和元数据。这可以看作是 __array__ 方法的反面。在每个实现 __array_wrap__ 的对象结束时，此方法将对具有最高 *数组优先级* 的输入对象调用，或者如果指定了输出对象，则对输出对象调用。 __array_priority__ 属性用于确定在存在多个返回对象的 Python 类型可能性的情况下，返回哪种类型的对象。例如，子类可以选择使用此方法将输出数组转换为子类的实例，并在将数组返回给用户之前更新元数据。

有关这些方法的更多信息，请参阅 ndarray 子类化和 ndarray 子类型化的特定功能。

互操作性示例#

示例：Pandas `Series` 对象#

考虑以下内容

>>> import pandas as pd
>>> ser = pd.Series([1, 2, 3, 4])
>>> type(ser)
pandas.core.series.Series

现在，ser **不是** ndarray，但由于它实现了 __array_ufunc__ 协议，我们可以像对待 ndarray 一样将 ufunc 应用于它

>>> np.exp(ser)
   0     2.718282
   1     7.389056
   2    20.085537
   3    54.598150
   dtype: float64
>>> np.sin(ser)
   0    0.841471
   1    0.909297
   2    0.141120
   3   -0.756802
   dtype: float64

我们甚至可以与其他 ndarray 进行操作

>>> np.add(ser, np.array([5, 6, 7, 8]))
   0     6
   1     8
   2    10
   3    12
   dtype: int64
>>> f(ser)
21.1977562209304
>>> result = ser.__array__()
>>> type(result)
numpy.ndarray

示例：PyTorch 张量#

PyTorch 是一个针对使用 GPU 和 CPU 的深度学习的优化张量库。PyTorch 数组通常称为 *张量*。张量类似于 NumPy 的 ndarray，不同之处在于张量可以在 GPU 或其他硬件加速器上运行。实际上，张量和 NumPy 数组通常可以共享相同的底层内存，从而无需复制数据。

>>> import torch
>>> data = [[1, 2],[3, 4]]
>>> x_np = np.array(data)
>>> x_tensor = torch.tensor(data)

请注意，x_np 和 x_tensor 是不同类型的对象

>>> x_np
array([[1, 2],
       [3, 4]])
>>> x_tensor
tensor([[1, 2],
        [3, 4]])

但是，我们可以将 PyTorch 张量视为 NumPy 数组，而无需显式转换

>>> np.exp(x_tensor)
tensor([[ 2.7183,  7.3891],
        [20.0855, 54.5982]], dtype=torch.float64)

此外，请注意此函数的返回类型与初始数据类型兼容。

警告

虽然这种 ndarray 和张量的混合可能很方便，但并不推荐这样做。它不适用于非 CPU 张量，并且在特殊情况下会产生意外的行为。用户应优先考虑将 ndarray 显式转换为张量。

注意

PyTorch 不实现 __array_function__ 或 __array_ufunc__。在后台，Tensor.__array__() 方法返回作为张量数据缓冲区的视图的 NumPy ndarray。有关详细信息，请参阅此问题和 __torch_function__ 实现。

还要注意，我们可以在此处看到 __array_wrap__ 的作用，即使 torch.Tensor 不是 ndarray 的子类

>>> import torch
>>> t = torch.arange(4)
>>> np.abs(t)
tensor([0, 1, 2, 3])

PyTorch 实现 __array_wrap__ 以便能够从 NumPy 函数中获取张量，并且我们可以直接修改它来控制从这些函数返回的哪种类型的对象。

示例：CuPy 数组#

CuPy 是一个与 NumPy/SciPy 兼容的数组库，用于使用 Python 进行 GPU 加速计算。CuPy 通过实现 cupy.ndarray（NumPy ndarray 的对应物）来实现 NumPy 接口的一个子集。

>>> import cupy as cp
>>> x_gpu = cp.array([1, 2, 3, 4])

cupy.ndarray 对象实现了 __array_ufunc__ 接口。这使得 NumPy ufunc 可以应用于 CuPy 数组（这将延迟操作到与 ufunc 相匹配的 CuPy CUDA/ROCm 实现）。

>>> np.mean(np.exp(x_gpu))
array(21.19775622)

请注意，这些操作的返回类型仍然与初始类型一致

>>> arr = cp.random.randn(1, 2, 3, 4).astype(cp.float32)
>>> result = np.sum(arr)
>>> print(type(result))
<class 'cupy._core.core.ndarray'>

有关详细信息，请参阅 CuPy 文档中的此页面。

cupy.ndarray 还实现了 __array_function__ 接口，这意味着可以执行以下操作

>>> a = np.random.randn(100, 100)
>>> a_gpu = cp.asarray(a)
>>> qr_gpu = np.linalg.qr(a_gpu)

CuPy 在 cupy.ndarray 对象上实现了许多 NumPy 函数，但并非所有函数都实现了。有关详细信息，请参阅 CuPy 文档。

示例：Dask 数组#

Dask 是一个用于 Python 中并行计算的灵活库。Dask 数组使用分块算法实现 NumPy ndarray 接口的一个子集，将大型数组分割成许多小数组。这允许使用多个内核对大于内存的数组进行计算。

Dask 支持 __array__() 和 __array_ufunc__。

>>> import dask.array as da
>>> x = da.random.normal(1, 0.1, size=(20, 20), chunks=(10, 10))
>>> np.mean(np.exp(x))
dask.array<mean_agg-aggregate, shape=(), dtype=float64, chunksize=(), chunktype=numpy.ndarray>
>>> np.mean(np.exp(x)).compute()
5.090097550553843

注意

Dask 是延迟评估的，在您通过调用 compute() 请求结果之前，不会计算计算结果。

有关详细信息，请参阅 Dask 数组文档和 Dask 数组与 NumPy 数组互操作性的范围。

示例：DLPack#

几个 Python 数据科学库实现了 __dlpack__ 协议。其中包括 PyTorch 和 CuPy。可以在 DLPack 文档的此页面上找到实现此协议的所有库的完整列表。

将 PyTorch CPU 张量转换为 NumPy 数组

>>> import torch
>>> x_torch = torch.arange(5)
>>> x_torch
tensor([0, 1, 2, 3, 4])
>>> x_np = np.from_dlpack(x_torch)
>>> x_np
array([0, 1, 2, 3, 4])
>>> # note that x_np is a view of x_torch
>>> x_torch[1] = 100
>>> x_torch
tensor([  0, 100,   2,   3,   4])
>>> x_np
array([  0, 100,   2,   3,   4])

导入的数组是只读的，因此写入或就地操作将失败

>>> x.flags.writeable
False
>>> x_np[1] = 1
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ValueError: assignment destination is read-only

必须创建副本才能就地对导入的数组进行操作，但这将意味着复制内存。不要对非常大的数组执行此操作

>>> x_np_copy = x_np.copy()
>>> x_np_copy.sort()  # works

注意

请注意，GPU 张量无法转换为 NumPy 数组，因为 NumPy 不支持 GPU 设备

>>> x_torch = torch.arange(5, device='cuda')
>>> np.from_dlpack(x_torch)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: Unsupported device in DLTensor.

但是，如果两个库都支持数据缓冲区所在的设备，则可以使用 __dlpack__ 协议（例如，PyTorch 和 CuPy）

>>> x_torch = torch.arange(5, device='cuda')
>>> x_cupy = cupy.from_dlpack(x_torch)

类似地，NumPy 数组可以转换为 PyTorch 张量

>>> x_np = np.arange(5)
>>> x_torch = torch.from_dlpack(x_np)

只读数组无法导出

>>> x_np = np.arange(5)
>>> x_np.flags.writeable = False
>>> torch.from_dlpack(x_np)  
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File ".../site-packages/torch/utils/dlpack.py", line 63, in from_dlpack
    dlpack = ext_tensor.__dlpack__()
TypeError: NumPy currently only supports dlpack for writeable arrays

进一步阅读#

数组接口协议
编写自定义数组容器
特殊属性和方法（有关 __array_ufunc__ 和 __array_function__ 协议的详细信息）
ndarray 子类化（有关 __array_wrap__ 和 __array_finalize__ 方法的详细信息）
ndarray 子类型化的特定功能（有关 __array_finalize__、__array_wrap__ 和 __array_priority__ 的实现的更多详细信息）
NumPy 路线图：互操作性
PyTorch 文档中关于与 NumPy 的桥梁的说明