NEP 31 — NumPy API 的情境局部與全域覆寫#

作者:

Hameer Abbasi <habbasi@quansight.com>

作者:

Ralf Gommers <rgommers@quansight.com>

作者:

Peter Bell <pbell@quansight.com>

狀態:

已取代

取代者:

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

類型:

標準追蹤

建立於:

2019-08-22

決議:

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

摘要#

此 NEP 提議透過可擴展的後端機制,使 NumPy 的所有公開 API 皆可覆寫。

接受此 NEP 意味著 NumPy 將在獨立的命名空間中提供全域和情境局部的覆寫,以及類似於 NEP-18 [2] 的調度機制。使用 __array_function__ 的初步經驗顯示,有必要能夠覆寫不接受類陣列引數的 NumPy 函數,因此無法透過 __array_function__ 覆寫。最迫切的需求是陣列建立和強制轉換函數,例如 numpy.zerosnumpy.asarray;請參閱例如 NEP-30 [9]

此 NEP 提議允許以選擇加入的方式,覆寫 NumPy API 的任何部分。它旨在作為 NEP-22 [3] 的全面解決方案,並避免為每個需要變得可覆寫的新函數或物件類型新增不斷增長的新協定列表。

動機與範疇#

此 NEP 的主要最終目標是使以下成為可能

# On the library side
import numpy.overridable as unp

def library_function(array):
    array = unp.asarray(array)
    # Code using unumpy as usual
    return array

# On the user side:
import numpy.overridable as unp
import uarray as ua
import dask.array as da

ua.register_backend(da) # Can be done within Dask itself

library_function(dask_array)  # works and returns dask_array

with unp.set_backend(da):
    library_function([1, 2, 3, 4])  # actually returns a Dask array.

在此,backend 可以是由 NumPy 或外部函式庫(例如 Dask 或 CuPy)定義的任何相容物件。理想情況下,它應該是模組 dask.arraycupy 本身。

這些覆寫類型對於終端使用者和函式庫作者都很有用。終端使用者可能已經編寫或希望編寫程式碼,然後稍後加速或移至不同的實作,例如 PyData/Sparse。他們可以透過簡單地設定後端來做到這一點。函式庫作者也可能希望編寫可在陣列實作之間移植的程式碼,例如 sklearn 可能希望為機器學習演算法編寫可在陣列實作之間移植的程式碼,同時也使用陣列建立函數。

此 NEP 採取整體方法:它假設 API 的某些部分需要可覆寫,並且這些部分會隨著時間推移而增長。它提供了一個通用框架和機制,以避免每次需要時都設計新的協定。這是 uarray 的目標:在 API 中允許覆寫,而無需設計新的協定。

此 NEP 提議以下事項:unumpy [8] 成為 NumPy API 中尚未被 __array_function____array_ufunc__ 涵蓋的部分的建議覆寫機制,並且 uarray 被供應商化到 NumPy 內的新命名空間中,以便使用者和下游依賴項可以存取這些覆寫。這種供應商化機制類似於 SciPy 決定為使 scipy.fft 可覆寫而採取的做法(請參閱 [10])。

uarray 背後的動機是多方面的:首先,已經進行了多次嘗試以允許調度 NumPy API 的部分,包括(最顯著地)NEP-13 [4] 中的 __array_ufunc__ 協定,以及 NEP-18 [2] 中的 __array_function__ 協定,但這已顯示需要開發進一步的協定,包括用於強制轉換的協定(請參閱 [5], [9])。參考文獻中已廣泛討論了需要這些覆寫的原因,並且此 NEP 不會嘗試詳細說明為什麼需要這些覆寫;但簡而言之:函式庫作者有必要能夠將任意物件強制轉換為他們自己類型的陣列,例如 CuPy 需要強制轉換為 CuPy 陣列,例如,而不是 NumPy 陣列。用更簡單的話來說,人們需要像 np.asarray(...) 這樣的東西或替代方案來「正常工作」並返回 duck 陣列。

使用方式與影響#

此 NEP 允許全域和情境局部的覆寫,以及類似 __array_function__ 的自動覆寫。

以下是此 NEP 將啟用的部分使用案例,除了動機章節中陳述的第一個案例之外

第一個是允許替代 dtype 返回其各自的陣列。

# Returns an XND array
x = unp.ones((5, 5), dtype=xnd_dtype) # Or torch dtype

第二個是允許覆寫 API 的部分。這是為了允許 np.linalg、BLAS 和 np.random 的替代和/或最佳化實作。

import numpy as np
import pyfftw # Or mkl_fft

# Makes pyfftw the default for FFT
np.set_global_backend(pyfftw)

# Uses pyfftw without monkeypatching
np.fft.fft(numpy_array)

with np.set_backend(pyfftw) # Or mkl_fft, or numpy
    # Uses the backend you specified
    np.fft.fft(numpy_array)

這將允許一種官方方式,使覆寫能夠與 NumPy 協同工作,而無需 monkeypatching 或發布 NumPy 的修改版本。

以下是一些其他使用案例,隱含但尚未陳述

data = da.from_zarr('myfile.zarr')
# result should still be dask, all things being equal
result = library_function(data)
result.to_zarr('output.zarr')

如果 magic_library 建構於 unumpy 之上,則第二個案例將會起作用。

from dask import array as da
from magic_library import pytorch_predict

data = da.from_zarr('myfile.zarr')
# normally here one would use e.g. data.map_overlap
result = pytorch_predict(data)
result.to_zarr('output.zarr')

有些後端可能依賴其他後端,例如 xarray 依賴 numpy.fft,並將時間軸轉換為頻率軸,或 Dask/xarray 持有 NumPy 陣列以外的陣列在其中。這將在程式碼內部以以下方式處理

with ua.set_backend(cupy), ua.set_backend(dask.array):
    # Code that has distributed GPU arrays here

向後相容性#

此 NEP 中未提議向後不相容的變更。

詳細描述#

提案#

此 NEP 在被接受時提出的唯一變更是使 unumpy 成為官方建議的覆寫 NumPy 的方式,以及透過 uarray 使某些子模組預設可覆寫。unumpy 將保持為獨立的儲存庫/套件(我們建議供應商化以避免硬依賴,並且僅在安裝時使用獨立的 unumpy 套件,而不是暫時依賴它)。具體而言,numpy.overridable 成為 unumpy 的別名,如果可用,則回退到供應商化的版本(如果不可用)。uarrayunumpy 將主要在 duck 陣列作者的輸入下開發,其次是自訂 dtype 作者,透過通常的 GitHub 工作流程。這樣做有幾個原因

  • 在發生錯誤或問題時,可以更快地迭代。

  • 在需要功能的情況下,可以更快地進行設計變更。

  • unumpy 也適用於舊版本的 NumPy。

  • 使用者和函式庫作者選擇加入覆寫流程,而不是在最不期望時發生中斷。簡而言之,unumpy 中的錯誤意味著 numpy 保持不受影響。

  • 對於 numpy.fftnumpy.linalgnumpy.random,主要命名空間中的函數將鏡像 numpy.overridable 命名空間中的函數。這樣做的原因是,即使對於 numpy.ndarray 輸入,這些子模組中也可能存在需要後端的函數。

unumpy 相較於其他解決方案的優勢#

相較於為遇到的每個問題定義新協定的方法,unumpy 提供了許多優勢:每當有需要覆寫的東西時,unumpy 都將能夠提供具有極小變更的統一 API。例如

  • ufunc 物件可以透過其 __call__reduce 和其他方法覆寫。

  • 其他函數可以以類似的方式覆寫。

  • np.asduckarray 消失,並變成設定後端的 np.overridable.asarray

  • 對於陣列建立函數(例如 np.zerosnp.empty 等)也是如此。

未來也是如此:使某物可覆寫僅需要對 unumpy 進行微小的變更。

unumpy 持有的另一個承諾是預設實作。可以根據其他 multimethod 為任何 multimethod 提供預設實作。這允許人們僅透過定義 NumPy API 的一小部分來覆寫其大部分。這是為了透過提供許多可以輕鬆地用其他函數表達的函數的預設實作,以及有助於實作大多數 duck 陣列將需要的功能的實用函數儲存庫,來簡化新 duck 陣列的建立。這將使我們能夠避免設計完整的協定,例如,用於堆疊和串聯的協定將被簡單地實作 stack 和/或 concatenate 來取代,然後為該類別中的其他所有內容提供預設實作。這同樣適用於轉置和許多其他尚未提出協定的函數,例如根據 in1disin,根據 uniquesetdiff1d 等。

它還允許人們以 __array_function__ 無法做到的方式覆寫函數,例如使用來自 opt_einsum 套件的版本覆寫 np.einsum,或 Intel MKL 覆寫 FFT、BLAS 或 ufunc 物件。他們將定義一個具有適當 multimethod 的後端,使用者將透過 with 語句選擇它們,或將它們註冊為後端。

最後一個好處是提供了一種清晰的方式來強制轉換為給定的後端(透過 ua.set_backend 中的 coerce 關鍵字),以及一個協定,用於不僅強制轉換陣列,還強制轉換 dtype 物件和 ufunc 物件以及來自其他函式庫的類似物件。這是因為實際存在第三方 dtype 套件,並且它們希望融入 NumPy 生態系統(請參閱 [6])。與 [7] 中提議的 C 級 dtype 重新設計相比,這是一個獨立的問題,它關於允許第三方 dtype 實作與 NumPy 協同工作,就像第三方陣列實作一樣。這些可以提供諸如單位、鋸齒狀陣列或其他 NumPy 範圍之外的功能。

在同一個檔案中混合 NumPy 和 unumpy#

通常,人們只想導入 unumpynumpy 中的一個,為了熟悉起見,您會將其導入為 np。但是,在某些情況下,人們可能希望混合 NumPy 和覆寫,並且有幾種方法可以做到這一點,具體取決於使用者的風格

from numpy import overridable as unp
import numpy as np

import numpy as np

# Use unumpy via np.overridable

Duck 陣列強制轉換#

numpy.arraynumpy.asarray 返回非 NumPy 陣列的物件存在固有的問題,尤其是在 C/C++ 或 Cython 程式碼的上下文中,這些程式碼可能會取得記憶體佈局與其預期不同的物件。但是,我們認為此問題不僅適用於這兩個函數,而且適用於所有返回 NumPy 陣列的函數。因此,使用者可以選擇加入覆寫,方法是使用子模組 numpy.overridable 而不是 numpy。NumPy 將繼續正常工作,不受 numpy.overridable 中任何內容的影響。

如果使用者希望取得 NumPy 陣列,則有兩種方法可以做到

  1. 使用 numpy.asarray(不可覆寫的版本)。

  2. 使用設定了 NumPy 後端並啟用強制轉換的 numpy.overridable.asarray

numpy.overridable 命名空間之外的別名#

numpy.randomnumpy.linalgnumpy.fft 中的所有功能都將別名到它們在 numpy.overridable 內各自的可覆寫版本。這樣做的原因是,存在 RNG(mkl-random)、線性代數例程(eigenblis)和 FFT 例程(mkl-fftpyFFTW)的替代實作,它們需要在 numpy.ndarray 輸入上運作,但仍然需要切換行為的能力。

這在幾個不同的方面與 monkeypatching 不同

  • 函數的呼叫者介面簽名始終相同,因此至少存在 API 合約的鬆散概念。Monkeypatching 不提供此能力。

  • 存在在本機切換後端的能力。

  • 已經有人建議,1.17 未登陸 Anaconda 預設通道的原因是 monkeypatching 和 __array_function__ 之間的不相容性,因為 monkeypatching 將完全繞過協定。

  • 形式為 from numpy import x; xnp.x 的語句將具有不同的結果,具體取決於導入是在 monkeypatching 發生之前還是之後進行的。

使用 __array_function____array_ufunc__ 完全不可能做到這一切。

已經正式認識到(至少部分地)為此需要後端系統,在 NumPy 路線圖中。

對於 numpy.random,仍然有必要使 C-API 符合 NEP-19 中提出的 API。對於 mkl-random 來說,這是不可行的,因為那樣就需要重寫它以適應該框架。關於流相容性的保證將與以前相同,但是如果存在影響 numpy.random 設定的後端,我們不對流相容性做出保證,並且由後端作者提供他們自己的保證。

提供隱式調度的方式#

有人建議需要調度不接受可調度方法的方法的能力,同時從另一個可調度方法中猜測後端。

作為一個具體範例,請考慮以下內容

with unumpy.determine_backend(array_like, np.ndarray):
    unumpy.arange(len(array_like))

雖然這在 uarray 中尚不存在,但新增它很簡單。需要這種程式碼的原因是,人們可能想要為提議的 *_like 函數或 like= 關鍵字引數提供替代方案。需要這些的原因是 NumPy API 中存在不接受可調度引數的函數,但仍然需要根據不同的可調度方法選擇後端。

需要選擇加入模組#

由於以下幾個原因,認識到需要選擇加入模組

  • API 的某些部分(例如 numpy.asarray)由於與 C/Cython 擴展的不相容性問題而根本無法覆寫,但是,人們可能希望使用設定了後端的 asarray 強制轉換為 duck 陣列。

  • 圍繞隱式選項和 monkeypatching 可能存在問題,例如上面提到的那些問題。

NEP 18 指出,這可能需要維護兩個單獨的 API。但是,例如,透過夾具單獨參數化 numpy.overridable 上的所有測試,可以減輕這種負擔。這也具有徹底測試它的副作用,不像 __array_function__。我們也認為,它提供了一個機會,可以將 NumPy API 合約與實作適當地分開。

終端使用者的好處以及混合後端#

uarray 中混合後端很容易,人們只需要執行

# Explicitly say which backends you want to mix
ua.register_backend(backend1)
ua.register_backend(backend2)
ua.register_backend(backend3)

# Freely use code that mixes backends here.

終端使用者的好處不僅僅限於編寫新程式碼。舊程式碼(通常以腳本形式)可以透過簡單的導入切換和一行新增首選後端,輕鬆移植到不同的後端。這樣,使用者可能會發現將現有程式碼移植到 GPU 或分散式運算更容易。

實作#

此 NEP 的實作將需要以下步驟

  • unumpy 儲存庫中實作對應於 NumPy API 的 uarray multimethod,包括用於覆寫 dtypeufuncarray 物件的類別,這些類別通常非常容易建立。

  • 將後端從 unumpy 移至各自的陣列函式庫。

透過參數化測試在 {numpy, unumpy} 上進行測試可以簡化維護。如果在方法中新增了新的引數,則需要在 unumpy 內更新對應的引數擷取器和替換器。

許多引數擷取器可以從 __array_function__ 協定的現有實作中重複使用,並且替換器通常可以在許多方法中重複使用。

對於預設可覆寫的命名空間部分,主要方法將需要重新命名並隱藏在 uarray multimethod 後面。

預設實作通常在文件中使用「等效於」等詞語看到,因此很容易獲得。

uarray 入門#

注意: 本節不會嘗試過於詳細地介紹 uarray,那是 uarray 文件化的目的。 [1] 但是,NumPy 社群將透過問題追蹤器為 uarray 的設計提供意見。

unumpy 介面定義了一組可覆寫的函式(多重方法),且與 NumPy API 相容。為此,它使用了 uarray 函式庫。uarray 是一個通用的工具,用於建立可分派到多個不同後端實作的多重方法。從這個意義上來說,它類似於 __array_function__ 協定,但關鍵區別在於後端是由終端使用者明確安裝的,而不是耦合到陣列類型中。

將後端與陣列類型解耦,為終端使用者和後端作者提供了更大的彈性。例如,可以:

  • 覆寫不將陣列作為引數的函式

  • 從陣列類型來源外部建立後端

  • 為同一陣列類型安裝多個後端

這種解耦也意味著 uarray 不受限於僅分派類似陣列的類型。後端可以自由檢查整個函式引數集,以確定它是否可以實作該函式,例如 dtype 參數分派。

定義後端#

uarray 由兩個主要協定組成:__ua_convert____ua_function__,依序調用,以及 __ua_domain____ua_convert__ 用於轉換和強制型轉。它的簽名是 (dispatchables, coerce),其中 dispatchablesua.Dispatchable 物件的可迭代物件,而 coerce 是一個布林值,指示是否強制轉換。ua.Dispatchable 是一個簡單的類別,由三個簡單的值組成:typevaluecoercible__ua_convert__ 傳回已轉換值的可迭代物件,或在失敗的情況下傳回 NotImplemented

__ua_function__ 的簽名是 (func, args, kwargs),並定義了函式的實際實作。它接收函式及其引數。如果傳回 NotImplemented,將導致移至函式的預設實作(如果存在),否則移至下一個後端。

以下是假設調用 uarray 多重方法時會發生的情況:

  1. 我們將引數標準化,因此任何沒有預設值的引數都放在 *args 中,而有預設值的引數則放在 **kwargs 中。

  2. 我們檢查後端列表。

    1. 如果列表為空,我們嘗試預設實作。

  3. 我們檢查後端的 __ua_convert__ 方法是否存在。如果存在:

    1. 我們將分派器的輸出傳遞給它,這是一個 ua.Dispatchable 物件的可迭代物件。

    2. 我們將此輸出連同引數一起饋送到引數替換器。NotImplemented 表示我們移至下一個後端的步驟 3。

    3. 我們將替換後的引數儲存為新引數。

  4. 我們將引數饋送到 __ua_function__ 中,並傳回輸出,如果輸出不是 NotImplemented,則退出。

  5. 如果預設實作存在,我們嘗試使用目前的後端來執行它。

  6. 如果失敗,我們移至下一個後端的步驟 3。如果沒有更多後端,我們移至步驟 7。

  7. 我們引發 ua.BackendNotImplementedError

定義可覆寫的多重方法#

要定義可覆寫的函式(多重方法),需要以下幾項:

  1. 一個分派器,傳回 ua.Dispatchable 物件的可迭代物件。

  2. 一個反向分派器,將可分派的值替換為提供的那些值。

  3. 一個域。

  4. 可選地,一個預設實作,可以用其他多重方法來提供。

作為範例,請考慮以下情況:

import uarray as ua

def full_argreplacer(args, kwargs, dispatchables):
    def full(shape, fill_value, dtype=None, order='C'):
        return (shape, fill_value), dict(
            dtype=dispatchables[0],
            order=order
        )

    return full(*args, **kwargs)

@ua.create_multimethod(full_argreplacer, domain="numpy")
def full(shape, fill_value, dtype=None, order='C'):
    return (ua.Dispatchable(dtype, np.dtype),)

unumpy 儲存庫中可以找到大量範例,[8]。這種覆寫可調用物件的簡單行為允許我們覆寫:

  • 方法

  • 屬性,透過 fgetfset

  • 整個物件,透過 __get__

NumPy 的範例#

實作類似 NumPy API 的函式庫將以下列方式使用它(作為範例):

import numpy.overridable as unp
_ua_implementations = {}

__ua_domain__ = "numpy"

def __ua_function__(func, args, kwargs):
    fn = _ua_implementations.get(func, None)
    return fn(*args, **kwargs) if fn is not None else NotImplemented

def implements(ua_func):
    def inner(func):
        _ua_implementations[ua_func] = func
        return func

    return inner

@implements(unp.asarray)
def asarray(a, dtype=None, order=None):
    # Code here
    # Either this method or __ua_convert__ must
    # return NotImplemented for unsupported types,
    # Or they shouldn't be marked as dispatchable.

# Provides a default implementation for ones and zeros.
@implements(unp.full)
def full(shape, fill_value, dtype=None, order='C'):
    # Code here

替代方案#

目前這個問題的替代方案是 NEP-18 [2]、NEP-13 [4] 和 NEP-30 [9] 的組合,再加上額外的協定(尚未指定)。即便如此,NumPy API 的某些部分仍將保持不可覆寫,因此這只是一個部分替代方案。

vendor unumpy 的主要替代方案是直接將其完全移入 NumPy,而不是作為單獨的套件發布。這也將實現提議的目標,但是我們目前傾向於將其保留為單獨的套件,原因已在上面說明。

第三個替代方案是將 unumpy 移入 NumPy 組織,並將其作為 NumPy 專案進行開發。這也將實現上述目標,並且也是本 NEP 可以考慮的可能性。但是,執行額外的 pip installconda install 可能會讓一些使用者不願意採用此方法。

需要選擇加入的替代方案主要是覆寫 np.asarraynp.array,並使 NumPy API 表面的其餘部分可覆寫,而是提供 np.duckarraynp.asduckarray 作為 duck-array 友善的替代方案,這些替代方案使用了各自的覆寫。但是,這樣做的缺點是會給 NumPy 調用增加少許額外負擔。

討論#

參考文獻和腳註#