NEP 35 — 陣列建立調度與 __array_function__#

作者:: Peter Andreas Entschev <pentschev@nvidia.com>
狀態:: 最終
類型:: 標準追蹤
建立於:: 2019-10-15
更新於:: 2020-11-06
決議:: https://mail.python.org/pipermail/numpy-discussion/2021-May/081761.html

摘要#

我們提議為所有陣列建立函數引入一個新的關鍵字引數 like=，以解決 NEP 18 [1] 所描述的 __array_function__ 的缺點之一。like= 關鍵字引數將建立引數型別的實例，從而可以直接建立非 NumPy 陣列。目標陣列型別必須實作 __array_function__ 協定。

動機與範疇#

許多函式庫實作了 NumPy API，例如用於圖形運算的 Dask、用於 GPGPU 運算的 CuPy、用於 N 維標籤陣列的 xarray 等。在底層，它們採用了 __array_function__ 協定，該協定允許 NumPy 理解和處理下游物件，如同它們是原生 numpy.ndarray 物件一樣。因此，社群在使用各種函式庫時，仍然受益於統一的 NumPy API。這不僅為標準化帶來了極大的便利，而且消除了學習新 API 和為每個新物件重寫程式碼的負擔。更技術性地說，協定的這種機制稱為「調度器」，這是我們從此以後在提及時使用的術語。

x = dask.array.arange(5)    # Creates dask.array
np.diff(x)                  # Returns dask.array

請注意上面我們如何透過呼叫 np.diff，經由 NumPy 命名空間呼叫 Dask 的 diff 實作，如果我們有 CuPy 陣列或任何其他來自採用 __array_function__ 的函式庫的陣列，情況也是如此。這允許編寫與實作函式庫無關的程式碼，因此使用者可以編寫一次程式碼，仍然能夠根據其需求使用不同的陣列實作。

顯然，如果陣列在其他地方建立並讓 NumPy 處理它們，那麼具備協定是很有用的。但這些陣列仍然必須在其原生函式庫中啟動並帶回。相反地，如果可以透過 NumPy API 建立這些物件，那麼將會有幾乎完整的體驗，全部使用 NumPy 語法。例如，假設我們有一些 CuPy 陣列 cp_arr，並想要一個類似的具有單位矩陣的 CuPy 陣列。我們仍然可以編寫以下內容

x = cupy.identity(3)

相反地，更好的方法是僅使用 NumPy API，現在可以使用以下方式實現

x = np.identity(3, like=cp_arr)

彷彿魔法一般，x 也會是 CuPy 陣列，因為 NumPy 能夠從 cp_arr 的型別推斷出來。請注意，如果沒有 like=，則最後一步是不可能的，因為 NumPy 不可能僅根據整數輸入就知道使用者期望的是 CuPy 陣列。

提議的新 like= 關鍵字僅旨在識別要調度的下游函式庫，並且該物件僅用作參考，這表示不會對該物件執行修改、複製或處理。

我們預期此功能對於函式庫開發人員最有用，允許他們根據使用者傳遞的陣列建立新的陣列以供內部使用，從而防止不必要地建立 NumPy 陣列，而這些陣列最終將額外轉換為下游陣列型別。

自 NumPy 1.17 以來，已放棄對 Python 2.7 的支援，因此我們使用 PEP-3102 [2] 中描述的僅限關鍵字引數標準來實作 like=，從而防止透過位置傳遞它。

用法與影響#

不使用來自下游函式庫的其他陣列的 NumPy 使用者可以繼續使用不帶 like= 引數的陣列建立常式。使用 like=np.ndarray 的效果就像沒有透過該引數傳遞任何陣列一樣。但是，這將產生額外的檢查，這將對效能產生負面影響。

為了理解 like= 的預期用途，在我們進入更複雜的情況之前，請考慮以下僅包含 NumPy 和 CuPy 陣列的說明性範例

import numpy as np
import cupy

def my_pad(arr, padding):
    padding = np.array(padding, like=arr)
    return np.concatenate((padding, arr, padding))

my_pad(np.arange(5), [-1, -1])    # Returns np.ndarray
my_pad(cupy.arange(5), [-1, -1])  # Returns cupy.core.core.ndarray

請注意上面的 my_pad 函數中，arr 如何用作參考，以指示填充應具有的陣列型別，然後再串連陣列以產生結果。另一方面，如果未使用 like=，NumPy 的情況仍然有效，但 CuPy 不允許這種自動轉換，最終會引發 TypeError: Only cupy arrays can be concatenated 例外狀況。

現在我們應該看看像 Dask 這樣的函式庫如何從 like= 中受益。在我們理解這一點之前，重要的是要稍微了解 Dask 的基礎知識以及它如何透過 __array_function__ 確保正確性。請注意，Dask 可以對不同種類的物件執行計算，例如資料框、包和陣列，這裡我們將嚴格關注陣列，這是我們可以與 __array_function__ 一起使用的物件。

Dask 使用圖形運算模型，這表示它將大型問題分解為許多較小的問題，並合併它們的結果以達到最終結果。為了將問題分解為更小的問題，Dask 也將陣列分解為更小的陣列，它稱之為「區塊」。因此，Dask 陣列可以由一個或多個區塊組成，並且它們可能具有不同的型別。但是，在 __array_function__ 的上下文中，Dask 僅允許相同型別的區塊；例如，Dask 陣列可以由多個 NumPy 陣列或多個 CuPy 陣列組成，但不能兩者混合。

為了避免計算期間的型別不符，Dask 在整個計算過程中將屬性 _meta 作為其陣列的一部分保留：此屬性用於在圖形建立時預測輸出型別，並建立某些函數計算中所需的任何中繼陣列。回到我們之前的範例，我們可以像下面看到的那樣，使用 _meta 資訊來識別我們將使用哪種陣列進行填充

import numpy as np
import cupy
import dask.array as da
from dask.array.utils import meta_from_array

def my_dask_pad(arr, padding):
    padding = np.array(padding, like=meta_from_array(arr))
    return np.concatenate((padding, arr, padding))

# Returns dask.array<concatenate, shape=(9,), dtype=int64, chunksize=(5,), chunktype=numpy.ndarray>
my_dask_pad(da.arange(5), [-1, -1])

# Returns dask.array<concatenate, shape=(9,), dtype=int64, chunksize=(5,), chunktype=cupy.ndarray>
my_dask_pad(da.from_array(cupy.arange(5)), [-1, -1])

請注意，上面傳回值中的 chunktype 如何從第一個 my_dask_pad 呼叫中的 numpy.ndarray 變更為第二個呼叫中的 cupy.ndarray。我們也在本範例中將函數重新命名為 my_dask_pad，目的是清楚地表明這將是 Dask 實作此類功能的方式（如果需要這樣做），因為它需要 Dask 的內部工具，這些工具在其他地方沒有太多用處。

為了啟用正確識別陣列型別，我們使用了 Dask 的實用函數 meta_from_array，它是作為支援 __array_function__ 的工作的一部分而引入的，允許 Dask 適當地處理 _meta。讀者可以將 meta_from_array 視為一個特殊函數，它僅傳回底層 Dask 陣列的型別，例如

np_arr = da.arange(5)
cp_arr = da.from_array(cupy.arange(5))

meta_from_array(np_arr)  # Returns a numpy.ndarray
meta_from_array(cp_arr)  # Returns a cupy.ndarray

由於 meta_from_array 傳回的值是類 NumPy 陣列，因此我們可以將其直接傳遞到 like= 引數中。

meta_from_array 函數主要針對函式庫的內部使用，以確保使用正確的型別建立區塊。如果沒有 like= 引數，將不可能確保 my_pad 建立的填充陣列的型別與輸入陣列的型別相符，這將導致 CuPy 引發 TypeError 例外狀況，如上文討論的 CuPy 單獨情況會發生的那樣。結合 Dask 對 meta 陣列的內部處理和提議的 like= 引數，現在可以處理涉及建立非 NumPy 陣列的情況，這可能是 Dask 目前從 __array_function__ 協定面臨的最嚴重的限制。

向後相容性#

此提案不會在 NumPy 內引發任何向後相容性問題，因為它僅為現有的陣列建立函數引入了一個新的關鍵字引數，其預設值為 None，因此不會變更目前的行為。

詳細描述#

__array_function__ 協定的引入允許下游函式庫開發人員使用 NumPy 作為調度 API。但是，該協定並未（也不打算）解決下游函式庫建立陣列的問題，從而阻止這些函式庫在該上下文中使用如此重要的功能。

此 NEP 的目的是以簡單明瞭的方式解決該缺點：引入一個新的 like= 關鍵字引數，類似於 empty_like 函數系列的工作方式。當陣列建立函數收到這樣的引數時，它們將觸發 __array_function__ 協定，並呼叫下游函式庫自己的陣列建立函數實作。like= 引數，顧名思義，應僅用於識別要調度的位置。與目前為止使用 __array_function__ 的方式（第一個引數識別目標下游函式庫）相反，並且為了避免破壞 NumPy 關於陣列建立的 API，新的 like= 關鍵字應僅用於調度目的。

下游函式庫將從 like= 引數中受益，而無需對其 API 進行任何變更，因為該引數僅需要由 NumPy 實作。仍然允許下游函式庫包含 like= 引數，因為它在某些情況下可能很有用，有關這些情況的詳細資訊，請參閱實作。仍然需要下游函式庫實作 __array_function__ 協定，如 NEP 18 [1] 所述，並適當地將引數引入其對 NumPy 陣列建立函數的呼叫中，如用法與影響中所例示。

實作#

實作需要為 NumPy 的所有現有陣列建立函數引入一個新的 like= 關鍵字。作為將新增此新引數的函數範例（但不限於），我們可以列舉那些採用類陣列物件的函數，例如 array 和 asarray、基於數值輸入建立陣列的函數（例如 range 和 identity），以及 empty 函數系列，即使這可能是多餘的，因為這些函數的特殊化版本已經存在，其命名格式為 empty_like。截至撰寫此 NEP 時，可以在 [5] 中找到陣列建立函數的完整清單。

此新提議的關鍵字應由 __array_function__ 機制從關鍵字字典中移除，然後再進行調度。這樣做的目的是雙重的

簡化那些已經選擇實作 __array_function__ 協定的函式庫對陣列建立的採用，從而消除了明確選擇所有陣列建立函數的要求；以及
大多數下游函式庫將不需要關鍵字引數，而那些需要的函式庫可以透過從 __array_function__ 擷取 self 來實現。

因此，下游函式庫不需要在其陣列建立 API 中包含 like= 關鍵字。在某些情況下（例如，Dask），具有 like= 關鍵字可能很有用，因為它將允許實作識別陣列內部結構。例如，Dask 可以從參考陣列中受益，以識別其區塊型別（例如，NumPy、CuPy、Sparse），從而建立由相同區塊型別支援的新 Dask 陣列，除非 Dask 可以讀取參考陣列的屬性，否則這是無法實現的。

函數調度#

有兩種不同的調度情況：Python 函數和 C 函數。為了允許 __array_function__ 調度，一種可能的實作是使用 overrides.array_function_dispatch 修飾 Python 函數，但 C 函數有不同的要求，我們將在稍後描述。

下面的範例顯示了關於如何使用 overrides.array_function_dispatch 修飾 asarray 的建議

def _asarray_decorator(a, dtype=None, order=None, *, like=None):
    return (like,)

@set_module('numpy')
@array_function_dispatch(_asarray_decorator)
def asarray(a, dtype=None, order=None, *, like=None):
    return array(a, dtype, copy=False, order=order)

請注意，在上面的範例中，實作保持不變，唯一的區別是修飾，它使用新的 _asarray_decorator 函數來指示 __array_function__ 協定在 like 不是 None 時進行調度。

我們現在將查看 C 函數範例，由於 asarray 無論如何都是 array 的特殊化版本，因此我們現在將以後者為例。由於 array 是 C 函數，因此目前 NumPy 對其 Python 原始碼所做的只是匯入該函數並調整其 __module__ 為 numpy。現在將使用 overrides.array_function_from_dispatcher 的特殊化版本修飾該函數，該版本也應負責調整模組。

array_function_nodocs_from_c_func_and_dispatcher = functools.partial(
    overrides.array_function_from_dispatcher,
    module='numpy', docs_from_dispatcher=False, verify=False)

@array_function_nodocs_from_c_func_and_dispatcher(_multiarray_umath.array)
def array(a, dtype=None, *, copy=True, order='K', subok=False, ndmin=0,
          like=None):
    return (like,)

上述 C 函數的實作有兩個缺點

它建立了另一個 Python 函數呼叫；以及
為了遵循目前的實作標準，文件應直接附加到 Python 原始碼。

此提案的第一個版本建議將上述實作作為 NumPy 中以 C 語言實作的函數的可行解決方案。但是，由於上面指出的缺點，我們已決定放棄 Python 端的任何變更，並透過純 C 語言實作解決這些問題。有關詳細資訊，請參閱 [7]。

在下游讀取參考陣列#

如實作區段開頭所述，like= 不會傳播到下游函式庫，但是，仍然可以存取它。這需要在下游函式庫的 __array_function__ 定義中進行一些變更，其中 self 屬性實際上是透過 like= 傳遞的。情況是這樣，因為我們使用 like= 作為調度陣列，這與 NEP-18 涵蓋的其他運算函數不同，後者通常在第一個位置引數上進行調度。

此類用途的一個範例是建立新的 Dask 陣列，同時保留其後端型別

# Returns dask.array<array, shape=(3,), dtype=int64, chunksize=(3,), chunktype=cupy.ndarray>
np.asarray([1, 2, 3], like=da.array(cp.array(())))

# Returns a cupy.ndarray
type(np.asarray([1, 2, 3], like=da.array(cp.array(()))).compute())

請注意，上面的陣列如何由 chunktype=cupy.ndarray 支援，並且計算後的結果陣列也是 cupy.ndarray。如果 Dask 沒有透過 __array_function__ 中的 self 屬性使用 like= 引數，則上面的範例將由 numpy.ndarray 支援

# Returns dask.array<array, shape=(3,), dtype=int64, chunksize=(3,), chunktype=numpy.ndarray>
np.asarray([1, 2, 3], like=da.array(cp.array(())))

# Returns a numpy.ndarray
type(np.asarray([1, 2, 3], like=da.array(cp.array(()))).compute())

鑑於函式庫將需要依賴 __array_function__ 中的 self 屬性來使用正確的參考陣列調度函數，我們建議以下兩種替代方案之一

在下游函式庫中引入支援 like= 引數的函數清單，並在呼叫函數時傳遞 like=self；或
檢查函數的簽章，並驗證它是否包含 like= 引數。請注意，這可能會導致更高的效能損失，並且假設可以進行內省，但如果函數是 C 函數，則可能無法進行內省。

為了使事情更清楚，讓我們看一下建議 2 如何在 Dask 中實作。Dask 中 __array_function__ 定義的目前相關部分如下所示

def __array_function__(self, func, types, args, kwargs):
    # Code not relevant for this example here

    # Dispatch ``da_func`` (da.asarray, for example) with *args and **kwargs
    da_func(*args, **kwargs)

這就是更新後的程式碼的外觀

def __array_function__(self, func, types, args, kwargs):
    # Code not relevant for this example here

    # Inspect ``da_func``'s  signature and store keyword-only arguments
    import inspect
    kwonlyargs = inspect.getfullargspec(da_func).kwonlyargs

    # If ``like`` is contained in ``da_func``'s signature, add ``like=self``
    # to the kwargs dictionary.
    if 'like' in kwonlyargs:
        kwargs['like'] = self

    # Dispatch ``da_func`` (da.asarray, for example) with args and kwargs.
    # Here, kwargs contain ``like=self`` if the function's signature does too.
    da_func(*args, **kwargs)

替代方案#

最近，NEP 37 [6] 提出了一個完全取代 __array_function__ 的新協定，這將需要已經採用 __array_function__ 的下游函式庫進行大量重工，因此我們仍然認為 like= 引數對於 NumPy 和下游函式庫是有益的。但是，該提案不一定被視為當前 NEP 的直接替代方案，因為它將完全取代 NEP 18，而此提案是建立在 NEP 18 之上的。關於此新提案的詳細資訊以及為什麼這將需要下游函式庫進行重工的討論，超出了當前提案的範圍。

討論#

參考文獻#

著作權#

本文檔已置於公共領域。