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 互操作的功能分為三組：

  1. 將外部物件轉換為 ndarray 的方法；

  2. 將執行從 NumPy 函數延遲到另一個陣列函式庫的方法；

  3. 使用 NumPy 函數並傳回外部物件實例的方法。

我們將在下方描述這些功能。

1. 在 NumPy 中使用任意物件#

NumPy API 中的第一組互操作性功能允許在可能的情況下將外部物件視為 NumPy 陣列。當 NumPy 函數遇到外部物件時，它們將嘗試（依序）：

  1. 緩衝區協定，在 Python C-API 文件 中描述。

  2. __array_interface__ 協定，在此頁面 中描述。作為 Python 緩衝區協定的前身，它定義了一種從其他 C 擴充功能存取 NumPy 陣列內容的方法。

  3. __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 維陣列物件的記憶體佈局。它為資料交換提供了以下語法：

  1. numpy.from_dlpack 函數，它接受具有 __dlpack__ 方法的（陣列）物件，並使用該方法建構一個包含來自 x 的資料的新陣列。

  2. 陣列物件上的 __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__：允許第三方物件支援和覆寫 ufuncs。

  • __array_function__：通用函數的 __array_ufunc__ 協定未涵蓋的 NumPy 功能的捕捉所有機制。

只要外部物件實作了 __array_ufunc__ 或 __array_function__ 協定，就可以在它們上操作，而無需顯式轉換。

__array_ufunc__ 協定#

通用函數（或簡稱 ufunc） 是函數的「向量化」包裝器，該函數接受固定數量的特定輸入並產生固定數量的特定輸出。如果並非所有輸入引數都是 ndarray，則 ufunc（及其方法）的輸出不一定是 ndarray。實際上，如果任何輸入定義了 __array_ufunc__ 方法，控制權將完全傳遞給該函數，即 ufunc 被覆寫。__array_ufunc__ 方法在該（非 ndarray）物件上定義，可以存取 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 子類別（子類型）的物件內部配置新陣列時，都會呼叫此方法。它可用於在建構後更改屬性，或從「父項」更新元資訊。

__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__ 實作 以瞭解詳細資訊。

另請注意，即使 torch.Tensor 不是 ndarray 的子類別，我們也可以在此處看到 __array_wrap__ 的作用

>>> 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 Array 使用分塊演算法實作 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

延伸閱讀#