NEP 37 — NumPy 類模組的調度協議

作者:

Stephan Hoyer <shoyer@google.com>

作者:

Hameer Abbasi

作者:

Sebastian Berg

狀態:

已取代

取代者:

NEP 56 — NumPy 主命名空間中的陣列 API 標準支援

類型:

標準追蹤

建立日期:

2019-12-29

決議:

https://mail.python.org/archives/list/numpy-discussion@python.org/message/Z6AA5CL47NHBNEPTFWYOTSUVSRDGHYPN/

摘要

NEP-18 的 __array_function__ 取得了好壞參半的成果。一些專案 (例如，dask、CuPy、xarray、sparse、Pint、MXNet) 熱情地採用了它。其他專案 (例如，JAX) 則比較不情願。在此，我們提出一個新的協議 __array_module__，我們預期最終可以取代 __array_function__ 的大多數用例。__array_module__ 協議需要使用者和函式庫作者明確採用，這確保了向後相容性，並且也比 __array_function__ 簡單得多，我們預期這兩者都將使其更容易採用。

為什麼 __array_function__ 還不夠

NEP-18 未能實現其目標的原因有兩個方面

1. 向後相容性考量。__array_function__ 對於使用它的函式庫具有重大影響

   • JAX 一直不願意實作 __array_function__，部分原因是它擔心破壞現有程式碼：使用者期望像 np.concatenate 這樣的 NumPy 函數返回 NumPy 陣列。這是 __array_function__ 設計的一個根本限制，我們選擇允許覆寫現有的 numpy 命名空間。像 Dask 和 CuPy 這樣的函式庫已經考慮並接受了 __array_function__ 的向後不相容性影響；如果這種影響不存在，對它們來說仍然會更好。

     請注意，像 PyTorch 和 scipy.sparse 這樣的專案也尚未採用 __array_function__，因為它們沒有與 NumPy 相容的 API 或語義。就 PyTorch 而言，這很可能在未來會添加。 scipy.sparse 的情況與 numpy.matrix 相同：它的語義與 numpy.ndarray 不相容，因此添加 __array_function__ (除非返回 NotImplemented 可能是個例外) 並不是一個健康的想法。

   • __array_function__ 目前需要「全有或全無」的方法來實作 NumPy 的 API。對於逐步採用沒有好的途徑，這對於已建立的專案來說尤其成問題，因為採用 __array_function__ 會導致重大變更。

2. 可以覆寫的內容的限制。 __array_function__ 有一些重要的缺口，最值得注意的是陣列建立和強制轉換函數

   • 陣列建立 常式 (例如，np.arange 和 np.random 中的常式) 需要一些其他機制來指示要建立哪種類型的陣列。NEP 35 提議為沒有現有陣列引數的函數新增可選的 like= 引數。但是，我們仍然缺乏任何機制來覆寫物件上的方法，例如 np.random.RandomState 所需的方法。

   • 陣列轉換 無法重複使用現有的強制轉換函數 (如 np.asarray)，因為 np.asarray 有時表示「轉換為精確的 np.ndarray」，而有時表示「轉換為_類似_ NumPy 陣列的東西」。這導致了 NEP 30 關於單獨的 np.duckarray 函數的提案，但這仍然無法解決如何將一個鴨子陣列轉換為與另一個鴨子陣列類型相符的類型。

其他被提出的可維護性問題包括

• 在支援覆寫的模組中，不再可能使用 NumPy 函數的別名。例如，CuPy 和 JAX 都設定 result_type = np.result_type，現在必須將 np.result_type 的使用包裝在它們自己的 result_type 函數中。

• 透過使用 NumPy 的實作來實作未實作的 NumPy 函數的 回退機制 很難正確完成 (但請參閱 dask 的版本)，因為 __array_function__ 沒有提供一致的介面。轉換陣列類型的所有引數需要遞迴到 *args, **kwargs 形式的通用引數中。

get_array_module 和 __array_module__ 協議

我們提出一個新的面向使用者的機制，用於調度到鴨子陣列實作，numpy.get_array_module。get_array_module 執行與 __array_function__ 相同的類型解析，並返回一個模組，該模組的 API 承諾與 numpy 的標準介面相符，可以實作對所有提供的陣列類型的操作。

該協議本身比 __array_function__ 更簡單且更強大，因為它不需要擔心實際實作函數。我們相信它解決了 __array_function__ 的大多數可維護性和功能限制。

新的協議是選擇加入、明確且具有本地控制；請參閱 附錄：API 覆寫的設計選擇，以討論這些設計特徵的重要性。

陣列模組合約

由 get_array_module/__array_module__ 返回的模組應盡最大努力在新陣列類型上實作 NumPy 的核心功能。未實作的功能應簡單地省略 (例如，存取未實作的函數應引發 AttributeError)。在未來，我們預期會編纂一個協議，用於請求 numpy 的受限子集；請參閱 請求 NumPy API 的受限子集 以取得更多詳細資訊。

如何使用 get_array_module

想要支援通用鴨子陣列的程式碼應明確呼叫 get_array_module，以確定要從哪個適當的陣列模組呼叫函數，而不是直接使用 numpy 命名空間。例如

# calls the appropriate version of np.something for x and y
module = np.get_array_module(x, y)
module.something(x, y)

陣列建立和陣列轉換都受到支援，因為調度由 get_array_module 而不是透過函數引數的類型來處理。例如，要使用隨機數產生函數或方法，我們可以簡單地取出適當的子模組

def duckarray_add_random(array):
    module = np.get_array_module(array)
    noise = module.random.randn(*array.shape)
    return array + noise

我們也可以從 NEP 30 撰寫鴨子陣列 stack 函數，而無需新的 np.duckarray 函數

def duckarray_stack(arrays):
    module = np.get_array_module(*arrays)
    arrays = [module.asarray(arr) for arr in arrays]
    shapes = {arr.shape for arr in arrays}
    if len(shapes) != 1:
        raise ValueError('all input arrays must have the same shape')
    expanded_arrays = [arr[module.newaxis, ...] for arr in arrays]
    return module.concatenate(expanded_arrays, axis=0)

預設情況下，如果沒有引數是陣列，get_array_module 將返回 numpy 模組。可以透過提供僅限關鍵字引數 module 來明確控制此回退。也可以透過設定 module=None 來指示應引發例外狀況，而不是返回預設陣列模組。

如何實作 __array_module__

實作想要支援 get_array_module 的鴨子陣列類型的函式庫需要實作相應的協議 __array_module__。這個新協議基於 Python 的算術調度協議，本質上是 __array_function__ 的簡化版本。

只有一個引數被傳遞到 __array_module__，即傳遞到 get_array_module 的唯一陣列類型的 Python 集合，也就是所有具有 __array_module__ 屬性的引數。

特殊方法應返回一個具有與 numpy 相符的 API 的命名空間，或 NotImplemented，表示它不知道如何處理該操作

class MyArray:
    def __array_module__(self, types):
        if not all(issubclass(t, MyArray) for t in types):
            return NotImplemented
        return my_array_module

從 __array_module__ 返回自訂物件

my_array_module 通常 (但不一定總是) 是一個 Python 模組。返回自訂物件 (例如，透過 __getattr__ 實作函數) 可能對某些進階用例很有用。

例如，自訂物件可以允許鴨子陣列模組的部分實作回退到 NumPy (儘管通常不建議這樣做，因為這種回退行為可能容易出錯)

class MyArray:
    def __array_module__(self, types):
        if all(issubclass(t, MyArray) for t in types):
            return ArrayModule()
        else:
            return NotImplemented

class ArrayModule:
    def __getattr__(self, name):
        import base_module
        return getattr(base_module, name, getattr(numpy, name))

從 numpy.ndarray 建立子類別

來自 NEP-18 的關於良好定義的類型轉換階層的所有相同指南仍然適用。numpy.ndarray 本身包含 __array_module__ 的匹配實作，這對於子類別來說很方便

class ndarray:
    def __array_module__(self, types):
        if all(issubclass(t, ndarray) for t in types):
            return numpy
        else:
            return NotImplemented

NumPy 的內部機制

get_array_module 的類型解析規則遵循與 Python 和 NumPy 現有的調度協議相同的模型：在超類別之前呼叫子類別，否則從左到右。__array_module__ 保證在每個唯一類型上僅被呼叫一次。

get_array_module 的實際實作將在 C 中，但應等效於此 Python 程式碼

def get_array_module(*arrays, default=numpy):
    implementing_arrays, types = _implementing_arrays_and_types(arrays)
    if not implementing_arrays and default is not None:
        return default
    for array in implementing_arrays:
        module = array.__array_module__(types)
        if module is not NotImplemented:
            return module
    raise TypeError("no common array module found")

def _implementing_arrays_and_types(relevant_arrays):
    types = []
    implementing_arrays = []
    for array in relevant_arrays:
        t = type(array)
        if t not in types and hasattr(t, '__array_module__'):
            types.append(t)
            # Subclasses before superclasses, otherwise left to right
            index = len(implementing_arrays)
            for i, old_array in enumerate(implementing_arrays):
                if issubclass(t, type(old_array)):
                    index = i
                    break
            implementing_arrays.insert(index, array)
    return implementing_arrays, types

與 __array_ufunc__ 和 __array_function__ 的關係

這些較舊的協議具有不同的用例，應保留

__array_module__ 旨在解決 __array_function__ 的限制，因此很自然地考慮它是否可以完全取代 __array_function__。這將提供雙重好處：(1) 簡化關於如何覆寫 NumPy 的使用者故事，以及 (2) 消除與在呼叫每個 NumPy 函數時檢查調度相關的效能降低。

但是，從使用者的角度來看，__array_module__ 和 __array_function__ 非常不同：它需要顯式呼叫 get_array_function，而不是簡單地重複使用原始的 numpy 函數。對於依賴鴨子陣列的函式庫來說，這可能還可以，但對於互動式使用來說，可能會非常冗長。

__array_ufunc__ 的一些調度用例也由 __array_module__ 解決，但並非全部。例如，仍然能夠以通用方式在非 NumPy 陣列 (例如，使用 dask.array) 上定義非 NumPy ufuncs (例如，來自 Numba 或 SciPy) 仍然很有用。

鑑於它們現有的採用和不同的用例，我們認為現在移除或棄用 __array_function__ 和 __array_ufunc__ 沒有意義。

用於實作 __array_function__ 和 __array_ufunc__ 的 Mixin 類別

儘管使用者介面有所不同，但 __array_module__ 和實作 NumPy API 的模組仍然包含使用現有鴨子陣列協議進行調度所需的足夠功能。

例如，以下 mixin 類別將根據 get_array_module 和 __array_module__ 為這些特殊方法提供合理的預設值

class ArrayUfuncFromModuleMixin:

    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
        arrays = inputs + kwargs.get('out', ())
        try:
            array_module = np.get_array_module(*arrays)
        except TypeError:
            return NotImplemented

        try:
            # Note this may have false positive matches, if ufunc.__name__
            # matches the name of a ufunc defined by NumPy. Unfortunately
            # there is no way to determine in which module a ufunc was
            # defined.
            new_ufunc = getattr(array_module, ufunc.__name__)
        except AttributeError:
            return NotImplemented

        try:
            callable = getattr(new_ufunc, method)
        except AttributeError:
            return NotImplemented

        return callable(*inputs, **kwargs)

class ArrayFunctionFromModuleMixin:

    def __array_function__(self, func, types, args, kwargs):
        array_module = self.__array_module__(types)
        if array_module is NotImplemented:
            return NotImplemented

        # Traverse submodules to find the appropriate function
        modules = func.__module__.split('.')
        assert modules[0] == 'numpy'
        for submodule in modules[1:]:
            module = getattr(module, submodule, None)
        new_func = getattr(module, func.__name__, None)
        if new_func is None:
            return NotImplemented

        return new_func(*args, **kwargs)

為了更容易撰寫鴨子陣列，我們也可以將這些 mixin 類別新增到 numpy.lib.mixins 中 (但上面的範例可能就足夠了)。

考慮的替代方案

命名

我們喜歡名稱 __array_module__，因為它反映了現有的 __array_function__ 和 __array_ufunc__ 協議。另一個合理的選擇可能是 __array_namespace__。

目前尚不清楚應該將呼叫此協議的 NumPy 函數稱為什麼 (get_array_module 在此提案中)。一些可能的替代方案：array_module、common_array_module、resolve_array_module、get_namespace、get_numpy、get_numpylike_module、get_duck_array_module。

請求 NumPy API 的受限子集

隨著時間的推移，NumPy 累積了非常大的 API 介面，僅在頂層 numpy 模組中就有超過 600 個屬性。任何鴨子陣列函式庫都不太可能或想要實作所有這些函數和類別，因為 NumPy 的常用子集要小得多。

我們認為定義 NumPy API 的「最小」子集 (省略很少使用或不建議使用的功能) 將會很有用。例如，最小 NumPy 可能包含 stack，但不包含其他堆疊函數 column_stack、dstack、hstack 和 vstack。這可以清楚地向鴨子陣列作者和使用者表明哪些功能是核心功能，以及他們可以跳過哪些功能。

支援請求 NumPy API 的受限子集將是包含在 get_array_function 和 __array_module__ 中的自然功能，例如，

# array_module is only guaranteed to contain "minimal" NumPy
array_module = np.get_array_module(*arrays, request='minimal')

為了方便使用 NumPy 進行測試以及與任何有效的鴨子陣列函式庫一起使用，當僅在 NumPy 陣列上呼叫 get_array_module 時，NumPy 本身將返回 numpy 模組的受限版本。省略的函數將根本不存在。

不幸的是，我們尚未弄清楚這些受限子集應該是什麼，因此現在這樣做沒有意義。當/如果我們這樣做時，我們可以將新的關鍵字引數新增到 get_array_module 或新增新的頂層函數，例如，get_minimal_array_module。我們還需要新增一個模仿 __array_module__ 的新協議 (例如，__array_module_minimal__)，或者可以將可選的第二個引數新增到 __array_module__ (使用 try/except 捕獲錯誤)。

用於隱式調度的新命名空間

與其在主要的 numpy 命名空間中使用 __array_function__ 支援覆寫，不如建立一個新的選擇加入命名空間，例如，numpy.api，其中包含支援調度的 NumPy 函數版本。這些覆寫將需要新的選擇加入協議，例如，模仿 __array_function__ 的 __array_function_api__。

這將透過選擇加入來解決 __array_function__ 的最大限制，並且還允許明確地覆寫像 asarray 這樣的函數，因為 np.api.asarray 將始終表示「轉換類陣列物件」。但它不會解決 __array_module__ 滿足的所有調度需求，並且會讓我們為了支援陣列使用者和實作人員而支援更複雜的協議。

我們可能會透過 __array_module__ 協議實作這樣一個新的命名空間。當然，一些使用者會發現這很方便，因為它需要的樣板程式碼稍微少一些。但這會讓使用者面臨一個令人困惑的選擇：他們應該在何時使用 get_array_module 與 np.api.something。此外，我們還必須新增和維護一個全新的模組，這比僅僅新增一個函數要昂貴得多。

在類型和陣列上而不是僅在類型上進行調度

與其僅透過唯一陣列類型支援調度，不如我們也可以透過陣列物件支援調度，例如，透過將 arrays 引數作為 __array_module__ 協議的一部分傳遞。這可能對具有元資料的陣列的調度很有用，例如 Dask 和 Pint 提供的元資料，但會對類型安全性和複雜性造成成本。

例如，一個同時支援 CPU 和 GPU 上陣列的函式庫可能會根據輸入引數決定在哪個裝置上從像 ones 這樣的函數建立新陣列

class Array:
    def __array_module__(self, types, arrays):
        useful_arrays = tuple(a in arrays if isinstance(a, Array))
        if not useful_arrays:
            return NotImplemented
        prefer_gpu = any(a.prefer_gpu for a in useful_arrays)
        return ArrayModule(prefer_gpu)

class ArrayModule:
    def __init__(self, prefer_gpu):
        self.prefer_gpu = prefer_gpu

    def __getattr__(self, name):
        import base_module
        base_func = getattr(base_module, name)
        return functools.partial(base_func, prefer_gpu=self.prefer_gpu)

這可能很有用，但不清楚我們是否真的需要它。Pint 似乎在沒有任何明確的陣列建立常式的情況下也能正常運作 (傾向於乘以單位，例如，np.ones(5) * ureg.m)，而且在大多數情況下，Dask 對於現有的 __array_function__ 樣式覆寫也還可以 (例如，傾向於使用 np.ones_like 而不是 np.ones)。選擇是否將陣列放置在 CPU 或 GPU 上可以透過 使陣列建立延遲 來解決。

附錄：API 覆寫的設計選擇

對於覆寫 NumPy 的 API，有許多可能的設計選擇。在此，我們討論了指導我們設計 __array_module__ 的設計決策的三個主要軸。

使用者的選擇加入與選擇退出

__array_ufunc__ 和 __array_function__ 協議提供了一種在 NumPy 現有命名空間內 覆寫 NumPy 函數的機制。這表示如果使用者不想要任何覆寫行為，則需要明確選擇退出，例如，透過使用 np.asarray() 轉換陣列。

從理論上講，這種方法降低了在使用者程式碼和函式庫中採用這些協議的門檻，因為使用標準 NumPy 命名空間的程式碼會自動相容。但實際上，這種方法並未奏效。例如，大多數維護良好的使用 NumPy 的函式庫都遵循使用 np.asarray() 轉換所有輸入的最佳實務，它們必須明確放寬此實務才能使用 __array_function__。我們的經驗是，使函式庫與新的鴨子陣列類型相容通常至少需要少量工作來適應資料模型和可以有效實作的操作中的差異。

這些選擇退出方法也大大複雜化了採用這些協議的函式庫的向後相容性，因為透過作為函式庫選擇加入，它們也會讓它們的使用者選擇加入，無論他們是否期望如此。為了贏得一直無法採用 __array_function__ 的函式庫，選擇加入方法似乎是必須的。

實作的明確與隱式選擇

__array_ufunc__ 和 __array_function__ 都對調度具有隱式控制：調度的函數是透過每個函數呼叫中的適當協議來確定的。這很好地概括了處理許多不同類型的物件，Python 中用於實作算術運算符號的用途證明了這一點，但它對可讀性有一個重要的缺點：程式碼的讀者不再立即清楚在呼叫函數時會發生什麼，因為函數的實作可能會被其任何引數覆寫。

速度 影響是

• 當使用鴨子陣列類型時，get_array_module 表示類型檢查只需要在每個支援鴨子類型的函數內部發生一次，而使用 __array_function__，它會在每次呼叫 NumPy 函數時發生。顯然，這將取決於函數，但如果典型的支援鴨子陣列的函數呼叫其他 NumPy 函數 3-5 次，則開銷會增加 3-5 倍。

• 當使用NumPy 陣列時，get_array_module 是每個函數的額外呼叫 ( __array_function__ 開銷保持不變)，這表示少量額外開銷。

顯式和隱式的實作選擇並非互斥的選項。實際上，我們熟悉的大多數透過 __array_function__ 覆寫 NumPy API 的實作（即 Dask、CuPy 和 Sparse，但不包括 Pint）也包含一種顯式方法，可以直接匯入模組來使用它們版本的 NumPy API（分別為 dask.array、cupy 或 sparse）。

區域、非區域與全域控制

最終的設計軸線是使用者如何控制 API 的選擇

• 區域控制，如多重分派和 Python 算術協定所例示，透過檢查類型或呼叫函數直接參數上的方法來決定要使用哪個實作。

• 非區域控制，例如 np.errstate，透過函數裝飾器或上下文管理器，使用全域狀態來覆寫行為。控制權以階層方式決定，透過最內層的上下文。

• 全域控制 提供一種機制，讓使用者設定預設行為，可以透過函數呼叫或設定檔。例如，matplotlib 允許設定繪圖後端的全域選擇。

區域控制通常被認為是 API 設計的最佳實踐，因為控制流程完全是顯式的，這使其最容易理解。非區域和全域控制偶爾會被使用，但通常是因為無知或缺乏更好的替代方案。

在 NumPy 公開 API 的鴨子型別情況下，我們認為非區域或全域控制會是錯誤，主要是因為它們不具備良好的組合性。如果一個函式庫設定/需要一組覆寫，然後在內部呼叫一個例程，該例程期望另一組覆寫，則最終的行為可能會非常令人驚訝。高階函數尤其成問題，因為函數被評估的上下文可能不是它們被定義的上下文。

我們認為非區域和全域控制適用的一類覆寫用例，是用於選擇一個保證具有完全一致介面的後端系統，例如 NumPy 陣列上 numpy.fft 的更快替代實作。然而，這些超出了目前提案的範圍，目前提案的重點是鴨子陣列。