NEP 18 — NumPy 高階陣列函數的調度機制#

作者:

Stephan Hoyer <shoyer@google.com>

作者:

Matthew Rocklin <mrocklin@gmail.com>

作者:

Marten van Kerkwijk <mhvk@astro.utoronto.ca>

作者:

Hameer Abbasi <hameerabbasi@yahoo.com>

作者:

Eric Wieser <wieser.eric@gmail.com>

狀態:

最終

類型:

標準追蹤

建立時間:

2018-05-29

更新時間:

2019-05-25

決議:

https://mail.python.org/pipermail/numpy-discussion/2018-August/078493.html

摘要#

我們提議 __array_function__ 協定,允許 NumPy 函數的參數定義該函數如何對它們運作。這將允許使用 NumPy 作為高效多維陣列運算的高階 API,即使陣列實作與 numpy.ndarray 大相逕庭。

詳細描述#

NumPy 的高階 ndarray API 已在 NumPy 本身之外針對不同架構實作多次,例如 GPU 陣列 (CuPy)、稀疏陣列 (scipy.sparse, pydata/sparse) 和平行陣列 (Dask array),以及深度學習框架中各種類似 NumPy 的實作,例如 TensorFlow 和 PyTorch。

同樣地,有許多專案建立在 NumPy API 之上,用於標記和索引陣列 (XArray)、自動微分 (Autograd, Tangent)、遮罩陣列 (numpy.ma)、物理單位 (astropy.units, pint, unyt) 等,這些專案在 NumPy API 之上新增了額外功能。這些專案大多也實作了 NumPy 高階 API 的近親變體。

我們希望能夠一起使用這些函式庫,例如,我們希望能夠將 CuPy 陣列放置在 XArray 中,或對 Dask 陣列程式碼執行自動微分。如果為 NumPy ndarray 撰寫的程式碼也可以被其他類似 NumPy 的專案使用,這將更容易實現。

例如,我們希望以下程式碼範例能夠與任何類似 NumPy 的陣列物件良好地協同運作

def f(x):
    y = np.tensordot(x, x.T)
    return np.mean(np.exp(y))

今天透過 NumPy 內部的各種協定機制,其中一部分是可能的。

  • np.exp 函數檢查 __array_ufunc__ 協定

  • .T 方法使用 Python 的方法調度運作

  • np.mean 函數明確檢查參數上是否有 .mean 方法

然而,其他函數,例如 np.tensordot,不會進行調度,而是可能強制轉換為 NumPy 陣列 (使用 __array__) 協定,或直接出錯。為了實現 NumPy API 的足夠覆蓋率以支援 XArray 和 autograd 等下游專案,我們希望支援 NumPy 內幾乎所有函數,這需要比僅 __array_ufunc__ 更廣泛的協定。我們希望有一個協定,允許 NumPy 函數的參數控制並將執行導向另一個函數 (例如 GPU 或平行實作),其方式在各專案之間是安全且一致的。

實作方式#

我們提議在 NumPy 中新增對新協定 __array_function__ 的支援。

此協定旨在作為通用函數 (如 np.exp) 的 __array_ufunc__ 協定未涵蓋的 NumPy 功能的包羅萬象的解決方案。語義與 __array_ufunc__ 非常相似,只是運算是由任意可呼叫物件而非 ufunc 實例和方法指定。

可以在 此筆記本 中找到原型實作。

警告

__array_function__ 協定及其在特定函數上的使用是實驗性的。我們計劃保留一個介面,使其可以覆寫 NumPy 函數,但是對於特定函數執行此操作的方式可能會且將會在幾乎沒有警告的情況下變更。如果無法接受此類降低的向後相容性保證,請勿依賴於非 NumPy 陣列的 NumPy 函數覆寫。請參閱下面的「非目標」以瞭解更多詳細資訊。

注意

使用 __array_function__ 協定的調度已實作,但尚未預設啟用

  • 在 NumPy 1.16 中,您需要在匯入 NumPy 之前設定環境變數 NUMPY_EXPERIMENTAL_ARRAY_FUNCTION=1,以測試 NumPy 函數覆寫。

  • 在 NumPy 1.17 中,協定將預設啟用,但可以使用 NUMPY_EXPERIMENTAL_ARRAY_FUNCTION=0 停用。

  • 最終,預期 __array_function__ 將永遠啟用。

介面#

我們為 __array_function__ 的實作提出以下簽名

def __array_function__(self, func, types, args, kwargs)
  • func 是 NumPy 公開 API 公開的任意可呼叫物件,以 func(*args, **kwargs) 的形式呼叫。

  • types 是來自原始 NumPy 函數呼叫的唯一引數類型 集合,這些類型實作了 __array_function__

  • 元組 args 和字典 kwargs 直接從原始呼叫傳遞。

__array_ufunc__ 不同,對於 func 的類型,或關於 argskwargs 中可能包含實作陣列 API 的物件,沒有高階保證。

為了方便 __array_function__ 實作者,types 提供了所有具有 '__array_function__' 屬性的引數類型。這允許實作者快速識別他們應該延遲到其他引數上的 __array_function__ 實作的情況。types 的類型有意模糊:frozenset 將最接近符合預期用途,但我們可能會基於效能原因而改用 tuple。在任何情況下,__array_function__ 實作都不應依賴 types 的迭代順序,這會違反明確定義的「類型轉換階層」(如 NEP-13 中所述)。

實作 NumPy API 的專案範例#

__array_function__ 的大多數實作將從兩個檢查開始

  1. 給定的函數是我們知道如何多載的函數嗎?

  2. 所有引數都是我們知道如何處理的類型嗎?

如果這些條件成立,__array_function__ 應傳回呼叫其針對 func(*args, **kwargs) 的實作結果。否則,它應傳回哨兵值 NotImplemented,表示這些類型未實作該函數。這比直接引發 TypeError 更佳,因為它讓其他引數有機會定義運算。

對於 __array_function__ 的傳回值沒有一般性要求,儘管大多數合理的實作可能應傳回與函數引數類型之一相同的陣列。如果/當 Python 獲得 協定的類型支援 且 NumPy 新增靜態類型註釋時,SupportsArrayFunction@overload 實作將指示 Any 的傳回類型。

定義自訂裝飾器 (implements,如下所示) 以註冊 __array_function__ 實作也可能很方便。

HANDLED_FUNCTIONS = {}

class MyArray:
    def __array_function__(self, func, types, args, kwargs):
        if func not in HANDLED_FUNCTIONS:
            return NotImplemented
        # Note: this allows subclasses that don't override
        # __array_function__ to handle MyArray objects
        if not all(issubclass(t, MyArray) for t in types):
            return NotImplemented
        return HANDLED_FUNCTIONS[func](*args, **kwargs)

def implements(numpy_function):
    """Register an __array_function__ implementation for MyArray objects."""
    def decorator(func):
        HANDLED_FUNCTIONS[numpy_function] = func
        return func
    return decorator

@implements(np.concatenate)
def concatenate(arrays, axis=0, out=None):
    ...  # implementation of concatenate for MyArray objects

@implements(np.broadcast_to)
def broadcast_to(array, shape):
    ...  # implementation of broadcast_to for MyArray objects

請注意,__array_function__ 實作不需要包含對應 NumPy 函數的所有可選引數 (例如,上面的 broadcast_to 省略了不相關的 subok 引數)。只有在 NumPy 函數呼叫中明確使用可選引數時,才會將它們傳遞給 __array_function__

注意

就像內建特殊方法 (如 __add__) 的情況一樣,正確撰寫的 __array_function__ 方法應始終在遇到未知類型時傳回 NotImplemented。否則,如果運算也包含您的物件之一,則將無法從另一個物件正確覆寫 NumPy 函數。

NumPy 程式碼庫本身內部的必要變更#

這將需要在 NumPy 程式碼庫內部進行兩項變更

  1. 一個函數,用於檢查可用的輸入,尋找這些輸入上的 __array_function__ 屬性,並適當地呼叫這些方法,直到其中一個成功。在常見的全 NumPy 案例中,這需要很快,並且即使在多載輸入數量很大時 (例如,np.concatenate 的情況),也需要具有可接受的效能 (不比線性時間差)。

    這是一個中等複雜度的額外函數。

  2. 在所有相關的 NumPy 函數中呼叫此函數。

    這會影響 NumPy 程式碼庫的許多部分,但複雜度非常低。

尋找並呼叫正確的 __array_function__#

給定一個 NumPy 函數、*args**kwargs 輸入,我們需要搜尋 *args**kwargs,以尋找可能具有 __array_function__ 屬性的所有適當輸入。然後,我們需要在這些可能的方法中進行選擇,並執行正確的方法。在幾個可能的實作之間協商可能很複雜。

尋找引數#

有效的引數可以直接在 *args**kwargs 中,例如在 np.tensordot(left, right, out=out) 的情況下,或者它們可能巢狀在清單或字典中,例如在 np.concatenate([x, y, z]) 的情況下。這可能會因兩個原因而產生問題

  1. 某些函數被賦予了長列表的值,而遍歷它們的成本可能過高。

  2. 某些函數可能具有我們不想檢查的引數,即使它們具有 __array_function__ 方法。

為了解決這些問題,NumPy 函數應明確指示它們的哪些引數可以多載,以及應如何檢查這些引數。作為規則,這應包括所有記錄為 array_likendarray 的引數。

我們建議透過為每個多載的 NumPy 函數編寫「調度器」函數來實現此目的

  • 這些函數將使用傳遞到 NumPy 函數的確切引數來呼叫 (即,dispatcher(*args, **kwargs)),並且應傳回要檢查覆寫的引數可迭代物件。

  • 調度器函數需要與其對應的 NumPy 函數共享完全相同的位置、可選和僅限關鍵字引數。否則,NumPy 函數的有效調用可能會在呼叫其調度器時導致錯誤。

  • 由於關鍵字引數的預設沒有 __array_function__ 屬性,因此依照慣例,我們將所有預設引數值設定為 None。這降低了簽名不同步的可能性,並最大限度地減少了調度器中的多餘資訊。唯一的例外應是引數值以某種方式影響調度的情況,這種情況應該很少見。

np.concatenate 的調度器範例可能具有啟發性

def _concatenate_dispatcher(arrays, axis=None, out=None):
    for array in arrays:
        yield array
    if out is not None:
        yield out

concatenate 調度器編寫為產生器函數,這使其能夠潛在包含可選 out 引數的值,而無需建立包含 (可能很長的) 要串聯的物件清單的新序列。

嘗試 __array_function__ 方法直到正確的方法生效#

許多引數可以實作 __array_function__ 協定。其中一些可能認為,在給定可用輸入的情況下,它們無法確定正確的結果。我們如何呼叫正確的方法?如果有多個有效方法,哪個具有優先權?

在大多數情況下,使用 __array_function__ 進行調度的規則與 __array_ufunc__ 的規則相符 (請參閱 NEP-13)。特別是

  • NumPy 將從所有指定的輸入中收集 __array_function__ 的實作,並按順序呼叫它們:子類別在超類別之前,否則從左到右。請注意,在某些涉及子類別的邊緣情況下,這與 Python 的 目前行為 略有不同。

  • __array_function__ 的實作透過傳回 NotImplemented 以外的任何值來指示它們可以處理運算。

  • 如果所有 __array_function__ 方法都傳回 NotImplemented,NumPy 將引發 TypeError

如果不存在 __array_function__ 方法,NumPy 將預設為呼叫其自己的實作,旨在用於 NumPy 陣列。例如,當所有類似陣列的引數都是 Python 數字或清單時,就會出現這種情況。(NumPy 陣列確實具有 __array_function__ 方法,如下所示,但如果 NumPy 陣列子類別以外的任何引數實作了 __array_function__,它始終會傳回 NotImplemented。)

__array_ufunc__ 的目前行為的一個偏差是,NumPy 將僅對每個唯一類型的第一個引數呼叫 __array_function__。這與 Python 的 呼叫反映方法的規則 相符,並且這確保了即使在有多個多載引數的情況下,檢查多載也具有可接受的效能。為了避免這兩個調度協定之間長期發散,我們也應該 更新 __array_ufunc__ 以符合此行為。

numpy.ndarray 上的 __array_function__ 方法#

具有 __array_function__ 的子類別的用例與具有 __array_ufunc__ 的子類別相同,因此 numpy.ndarray 也定義了 __array_function__ 方法

def __array_function__(self, func, types, args, kwargs):
    if not all(issubclass(t, ndarray) for t in types):
        # Defer to any non-subclasses that implement __array_function__
        return NotImplemented

    # Use NumPy's private implementation without __array_function__
    # dispatching
    return func._implementation(*args, **kwargs)

此方法與 NumPy 的調度規則相符,因此在大多數情況下,可以假裝 ndarray.__array_function__ 不存在。在 array_function_dispatch 裝飾器中定義的私有 _implementation 屬性允許我們避免 __array_ufunc__ 協定中需要的 NumPy 陣列的特殊情況。

__array_function__ 協定始終在超類別之前呼叫子類別,因此如果任何 ndarray 子類別參與運算,它們將有機會覆寫它,就像任何其他引數覆寫 __array_function__ 一樣。但是,組合基礎 NumPy 陣列和子類別的運算中的預設行為是不同的:如果子類別傳回 NotImplemented,則將呼叫 NumPy 的函數實作,而不是引發例外。這是適當的,因為預期子類別是 可替換的

我們仍然告誡子類別的作者,在依賴 NumPy 內部實作的詳細資訊時要謹慎。並非總是能夠編寫完全可替換的 ndarray 子類別,例如,在涉及建立新陣列的情況下,尤其因為 NumPy 利用專門針對基礎 NumPy 陣列的內部優化,例如以 C 語言編寫的程式碼。即使 NumPy 的實作今天碰巧有效,將來也可能無效。在這些情況下,您的補救措施是透過子類別上的 __array_function__ 重新實作頂層 NumPy 函數。

NumPy 函數內部的變更#

給定一個定義上述行為的函數,暫時稱其為 implement_array_function,我們現在需要從每個相關的 NumPy 函數內部呼叫該函數。這是一個普遍性的變更,但程式碼相當簡單且無害,如果沒有任何引數實作 __array_function__ 協定,則應快速完成且沒有影響。

為了實現此目的,我們定義了一個 array_function_dispatch 裝飾器來重寫 NumPy 函數。基本實作如下

def array_function_dispatch(dispatcher, module=None):
    """Wrap a function for dispatch with the __array_function__ protocol."""
    def decorator(implementation):
        @functools.wraps(implementation)
        def public_api(*args, **kwargs):
            relevant_args = dispatcher(*args, **kwargs)
            return implement_array_function(
                implementation, public_api, relevant_args, args, kwargs)
        if module is not None:
            public_api.__module__ = module
        # for ndarray.__array_function__
        public_api._implementation = implementation
        return public_api
    return decorator

# example usage
def _broadcast_to_dispatcher(array, shape, subok=None):
    return (array,)

@array_function_dispatch(_broadcast_to_dispatcher, module='numpy')
def broadcast_to(array, shape, subok=False):
    ...  # existing definition of np.broadcast_to

使用裝飾器非常棒!我們不需要變更現有 NumPy 函數的定義,只需要為調度器函數編寫幾行額外程式碼。我們甚至可以為具有相同簽名 (例如,sumprod) 的函數族重複使用單個調度器。對於此類函數,最大的變更可能是向 docstring 新增幾行,以註記哪些引數已檢查多載。

特別值得一提的是裝飾器對 functools.wraps 的使用

  • 這確保了包裝函數具有與包裝的 NumPy 函數相同的名稱和 docstring。

  • 在 Python 3 上,它還確保裝飾器函數複製原始函數簽名,這對於基於內省的工具 (例如自動完成) 非常重要。

  • 最後,它確保包裝函數 可以被 pickled

範例用法說明了 NumPy 貢獻者相關的幾個撰寫調度器的最佳實務

  • 我們傳遞了 module 引數,它反過來設定了產生函數上的 __module__ 屬性。這是為了獲得更好的錯誤訊息的好處,這裡用於 NumPy 內部在找不到實作時引發的錯誤,例如,TypeError: no implementation found for 'numpy.broadcast_to'。將 __module__ 設定為 NumPy 公開 API 中的規範位置鼓勵使用者使用 NumPy 的公開 API 來識別 __array_function__ 中的函數。

  • 調度器是一個傳回元組的函數,而不是使用 yield 的等效 (且同樣有效) 產生器

    # example usage
    def broadcast_to(array, shape, subok=None):
        yield array
    

    這絕非偶然:當調度器函數傳回內建序列類型 (tuplelist) 時,NumPy 的 __array_function__ 調度實作速度最快。

    在相關的注意事項中,即使在某些情況下您知道它們不能具有 __array_function__ 方法,調度器傳回引數也完全可以。這可能會針對具有預設引數 (例如,None) 或複雜簽名的函數發生。NumPy 的調度邏輯可以非常快速地解決這些情況,因此通常不值得自己解析它們。

注意

上述 array_function_dispatch 的程式碼已從此 NEP 的原始版本更新,以符合 NumPy 中的實際實作

擴展性#

這種方法的一個重要優點是,它允許在不破壞已依賴 __array_function__ 的程式碼的情況下,向 NumPy 函數添加新的可選參數。

這不是理論上的問題。NumPy 較舊、隨意地在 np.sum() 等函數內部實作覆寫,當我們決定添加新的可選參數時,例如,新的 keepdims 參數僅在使用的情況下傳遞,這就需要一些笨拙的變通方法。

def sum(array, ..., keepdims=np._NoValue):
    kwargs = {}
    if keepdims is not np._NoValue:
        kwargs['keepdims'] = keepdims
    return array.sum(..., **kwargs)

對於 __array_function__ 的實作者來說,這也意味著即使是現有的可選參數也可以逐步實作,並且僅在有意義的情況下實作。例如,實作不可變陣列的函式庫不需要明確地在函數簽名中包含不受支援的 out 參數。這可能有點難以正確實作,例如,

def my_sum(array, ..., out=None):
    if out is not None:
        raise TypeError('out argument is not supported')
    ...

因此,我們避免鼓勵在 NumPy 呼叫的函數簽名中添加通用的 **ignored_kwargs 這個誘人的捷徑,因為這會對拼寫錯誤或忽略的參數靜默失敗。

效能#

效能始終是 NumPy 關注的問題,即使 NumPy 使用者在選擇 Python 語言本身時,已經將可用性優先於純粹的速度。重要的是,這種新的 __array_function__ 協定不會在 NumPy 函數作用於 NumPy 陣列的典型情況下,造成顯著的成本。

我們的微基準測試結果顯示,上述覆寫機制的純 Python 實作,在每次 NumPy 函數呼叫中,在沒有任何重載參數的情況下,大約增加了 2-3 微秒的額外開銷。作為參考,小型陣列上的典型 NumPy 函數的執行時間為 1-10 微秒,這主要取決於函數邏輯中有多少部分是用 C 語言編寫的。例如,一微秒大約是 ndarray.sum() 方法 (1.6 微秒) 和 numpy.sum() 函數 (2.6 微秒) 之間的速度差異。

幸運的是,我們預期 implement_array_function 的 C 實作會顯著降低額外開銷,這也是執行時間的主要部分。這將使 array_function_dispatch 裝飾器和分派器函數本身增加約 0.5 微秒的額外開銷,在典型情況下,總共可能約為 1 微秒的額外開銷。

我們認為,對於用 Python 編寫的程式碼來說,這種程度的額外開銷是可以接受的。我們非常確定,絕大多數 NumPy 使用者並不關心 NumPy 函數上以微秒為單位的效能差異,因為在 Python 中,很難在不到一微秒的時間內完成任何事情

在 NumPy 之外使用#

此協定沒有任何特定於 NumPy 本身的東西。我們是否應該鼓勵第三方函式庫使用相同的 __array_function__ 協定來重載非 NumPy 函數,例如,在 SciPy 中實現陣列實作的通用功能?

這將提供顯著的優勢 (SciPy 不需要發明自己的分派系統),並且我們認為沒有任何缺點,因為每個使用 __array_function__ 進行分派的函數都需要被明確識別。像 Dask、CuPy 和 Autograd 這樣的函式庫已經以類似於它們包裝 NumPy 的方式,包裝了 SciPy 功能的有限子集 (例如,scipy.linalg)。

如果我們想這樣做,我們至少應該公開裝飾器 array_function_dispatch(),並且可能也公開較低層級的 implement_array_function(),作為 NumPy 公共 API 的一部分。

非目標#

我們的目標是基本策略,可以在相對較短的時間內,即單個 NumPy 版本的開發週期內,機械式地應用於 NumPy API 中的幾乎所有函數。

我們希望在第一次嘗試時就正確地完成 __array_function__ 協定和所有特定的重載,但我們在這裡的明確目標是獲得一些大部分可以運作的東西 (並且可以迭代改進),而不是等待最佳實作。快速行動的代價是,目前此協定應被視為嚴格的實驗性質。我們保留在未來隨時更改此協定的細節以及特定 NumPy 函數如何使用它的權利 – 即使在其他方面僅為錯誤修復的 NumPy 版本中也是如此。實際上,一旦 __array_function__ 的初始問題解決,我們將使用縮短的棄用週期,短至單個主要 NumPy 版本 (例如,最少四個月)。

特別是,我們不打算編寫額外的 NEP 來列出所有要重載的特定函數,以及它們應該如何重載。我們將把這留給個人提取請求的提交者自行決定,並相信他們會將任何爭議提交給相關方討論。

但是,我們已經知道幾個函數族應該明確排除在 __array_function__ 之外。這些將需要它們自己的協定

  • 通用函數,它們已經有自己的協定。

  • arrayasarray,因為它們明確旨在強制轉換為實際的 numpy.ndarray 物件。

  • 任何類型的方法的分派,例如,np.random.RandomState 物件上的方法。

我們也預期,最初將使用 __array_function__ 協定的特定函數的覆寫機制,在未來可能會改變。作為我們預期未來會破壞行為的一個具體範例,某些函數 (例如 np.where) 目前不是 NumPy 通用函數,但可以想像在未來可能會變成通用函數。當/如果這種情況發生時,我們將把此類重載從使用 __array_function__ 更改為更專業的 __array_ufunc__

向後相容性#

此提案不會更改現有的語義,但目前具有 __array_function__ 屬性的參數除外,這種情況應該很少見。

替代方案#

專門協定#

我們可以 (並且應該) 繼續開發像 __array_ufunc__ 這樣的協定,用於 NumPy 功能的內聚子集。

如上所述,如果這意味著我們用 __array_function__ 重載的某些函數應該切換到新的協定,只要 __array_function__ 保持其實驗性狀態,這就絕對可以接受。

切換到新的協定應該使用 NumPy 正常棄用週期的縮短版本

  • 對於單個主要版本,在檢查任何新協定之後,NumPy 仍然應該檢查 __array_function__ 方法,這些方法實作給定的函數。如果任何參數從 __array_function__ 返回的值不是 NotImplemented,則應發出描述性的 FutureWarning

  • 在下一個主要版本中,將移除對 __array_function__ 的檢查。

單獨命名空間#

用於重載函數的單獨命名空間是另一種可能性,無論是在 NumPy 內部還是外部。

這具有減輕任何可能關於向後相容性的擔憂的優點,並將為快速實驗提供最大的自由度。從長遠來看,它將提供一個乾淨的抽象層,將 NumPy 的高階 API 與 numpy.ndarray 物件上的預設實作分開。

缺點是這將需要所有現有程式碼的明確選擇加入,例如,import numpy.api as np,並且從長遠來看,將導致維護兩個單獨的 NumPy API。此外,numpy 本身的許多函數已經被重載 (但不足),因此關於 NumPy 中高階與低階 API 的混淆仍然存在。

或者,可以為 NumPy 高階 API 的非重載版本創建一個單獨的命名空間,例如 numpy.array_only,以便在對 NumPy 陣列的效能至關重要的情況下使用。這與單獨的命名空間具有大多數相同的缺點。

多重分派#

我們提出的 __array_function__ 協定的替代方案是將 NumPy 的核心函數實作為多重方法。儘管我們中的一位編寫了一個用於 Python 的 多重分派函式庫,但我們認為這種方法在近期內對於 NumPy 來說沒有意義。

主要原因是 NumPy 已經有一個經過驗證的分派機制 __array_ufunc__,它基於 Python 自己的算術分派系統,並且添加另一個以非常不同的方式運作的機制會令人困惑。這也將是對 NumPy 本身更具侵入性的更改,NumPy 需要獲得多重分派實作。

在未來,NumPy 高階 API 的多重分派實作可能是合理的。幸運的是,__array_function__ 並不排除這種可能性,因為用多重分派來編寫預設 __array_function__ 實作的 shim 將是直接了當的。

以有限的核心 API 實作#

某些 NumPy 函數的內部實作非常簡單。例如

  • np.stack() 僅透過將索引與 np.newaxisnp.concatenateshape 屬性結合,以幾行程式碼實作。

  • np.mean() 在內部以 np.sum()np.divide().astype().shape 實作。

這暗示了定義最小「核心」ndarray 介面的可能性,並在 NumPy 內部依賴它來實作完整的 API。這是一個有吸引力的選項,因為它可以顯著減少新陣列實作所需的工作。

但是,這也帶來了幾個缺點

  1. NumPy 如何根據重載函數實作高階函數的細節,現在成為 NumPy 公共 API 的隱含部分。例如,將 stack 重構為在內部使用 np.block() 而不是 np.concatenate(),現在將成為一個破壞性變更。

  2. 陣列函式庫可能更喜歡以不同於 NumPy 的方式實作高階函數。例如,函式庫可能更喜歡直接實作像 mean() 這樣的基本運算,而不是依賴 sum() 後跟除法。更一般地說,目前尚不清楚到底什麼符合核心功能,而弄清楚這一點可能是一個大型專案。

  3. 我們還沒有用於重載陣列物件上的屬性和方法的系統,例如,用於存取 .dtype.shape。這應該是未來 NEP 的主題,但在那之前,我們應該不願意依賴這些屬性。

鑑於這些擔憂,我們認為支援幾乎 NumPy API 中每個公共函數的顯式重載是有價值的。這並不排除未來使用簡化的核心功能和 __array_function__ 以及用於確保陣列公開像 numpy.ndarray 這樣的方法和屬性的協定和/或基底類別來重寫 NumPy 函數的未來可能性。但是,為了良好地運作,這將需要實作某些但不是所有具有 __array_function__ 的函數的可能性,例如,在下一節中描述的那樣。

NumPy API 的部分實作#

在目前的設計中,實作 __array_function__ 以重載至少一個函數的類別,隱含地宣告了實作整個 NumPy API 的意圖。不可能僅在類型上實作 np.concatenate(),但對於所有其他函數,退回到使用 np.asarray() 進行轉換的 NumPy 預設行為。

這可能會帶來向後相容性問題,這將阻止函式庫以增量方式採用 __array_function__。例如,目前大多數 numpy 函數會隱含地將 pandas.Series 物件轉換為 NumPy 陣列,這種行為肯定許多 pandas 使用者都依賴。如果 pandas 僅針對 np.concatenate 實作 __array_function__,則不相關的 NumPy 函數 (例如 np.nanmean) 會突然在 pandas 物件上中斷,並引發 TypeError。

即使是重新實作大多數 NumPy 公共 API 的函式庫,有時也會依賴於使用 NumPy 中的實用函數,而無需包裝器。例如,CuPy 和 JAX 都只是 使用別名np.result_type,它已經支援具有 dtype 屬性的 duck-type。

使用 __array_ufunc__,可以透過將所有參數轉換為 numpy 陣列並重新呼叫 ufunc 來減輕這種擔憂,但 __array_function__ 支援的異質函數簽名使得不可能為 __array_function__ 實作這種通用的回退行為。

我們考慮了三種可能的方法來解決這個問題,但沒有一種是完全令人滿意的

  1. 更改從 __array_function__ 返回 NotImplemented 的所有參數的含義,以指示所有參數都應強制轉換為 NumPy 陣列,並且應重試操作。但是,許多陣列函式庫 (例如 scipy.sparse) 真的不想要隱含轉換為 NumPy 陣列,並且通常會避免實作 __array__,正是因為這個原因。隱含轉換可能會導致靜默錯誤和效能降低。

    潛在地,我們可以僅針對實作 __array__ 的類型啟用此行為,這將解決像 scipy.sparse 這樣最成問題的情況。但在實踐中,很大一部分呈現像 NumPy 陣列這樣的高階 API 的類別已經實作了 __array__。這將排除在這些物件上可靠地使用 NumPy 的高階 API。

  2. 使用某種其他 sentinel 值,例如 np.NotImplementedButCoercible,來指示實作部分 NumPy 更高階陣列 API 的類別,在回退時是可強制轉換的。如果所有參數都返回 NotImplementedButCoercible,則會強制轉換參數,並重試操作。

    不幸的是,在遇到 NotImplementedButCoercible 後的正確行為並不總是顯而易見的。特別具有挑戰性的是「混合」情況,其中某些參數返回 NotImplementedButCoercible,而其他參數返回 NotImplemented。在僅強制轉換「可強制轉換」的參數後,分派是否會重試?如果是這樣,那麼可以想像我們會無限次地循環遍歷分派邏輯。無論哪種方式,分派規則肯定會變得更複雜且更難以理解。

  3. 允許存取 NumPy 的函數實作,例如,以公開的 __skip_array_function__ 屬性的形式在 NumPy 函數上。這將允許透過在 __array_function__ 方法內部使用 func.__skip_array_function__ 來回退到 NumPy 的實作,並且也可能用於避免分派的額外開銷。但是,它有暴露 NumPy 函數的 NumPy 實作細節的風險,這些 NumPy 函數不會在內部呼叫 np.asarray()。請參閱 此註釋 以獲取完整討論的摘要。

這些解決方案將解決實際的用例,但代價是增加了額外的複雜性。我們希望在使用 __array_function__ 之前獲得經驗,瞭解它實際上是如何使用的,然後再做出難以撤銷的決定。

檢查類型註釋的神奇裝飾器#

原則上,Python 3 類型註釋包含足夠的資訊來自動創建大多數 dispatcher 函數。使用這些註釋來免除手動編寫分派器的需要將很方便,例如,

@array_function_dispatch
def broadcast_to(array: ArrayLike
                 shape: Tuple[int, ...],
                 subok: bool = False):
    ...  # existing definition of np.broadcast_to

這將需要某種形式的自動程式碼產生,無論是在編譯時還是匯入時。

我們認為這是一個有趣的可能擴展,可以在未來考慮。我們認為現在這樣做沒有意義,因為程式碼產生涉及權衡,而 NumPy 在類型註釋方面的經驗仍然非常有限。即使 NumPy 僅適用於 Python 3 (這將在 2019 年的某個時候 發生),我們還沒有準備好直接註釋 NumPy 的程式碼庫。

支援實作特定的參數#

我們可以允許 __array_function__ 實作透過在分派器函數中包含 **ignored_kwargs 來添加它們自己的可選關鍵字參數,例如,

def _concatenate_dispatcher(arrays, axis=None, out=None, **ignored_kwargs):
    ...  # same implementation of _concatenate_dispatcher as above

實作特定的參數在其他方面模擬 NumPy 更高階 API 的函式庫中相當常見 (例如,dask.array.sum() 添加了 split_every,而 tensorflow.reduce_sum() 添加了 name)。在 NumPy 中支援它們對於在 NumPy 函數之上實作新的高階陣列函數的函式庫特別有用,例如,

def mean_squared_error(x, y, **kwargs):
    return np.mean((x - y) ** 2, **kwargs)

否則,我們將需要針對每個陣列實作的不同版本 mean_squared_error,以便將實作特定的參數傳遞給 mean()

我們不允許添加可選的位置參數,因為這些參數是為 NumPy 本身將來使用而保留的,但關鍵字參數之間的衝突應該相對較少。

但是,這種靈活性會帶來成本。特別是,它隱含地將 **kwargs 添加到所有包裝的 NumPy 函數的簽名中,而實際上並沒有包含它 (因為我們使用 functools.wraps)。這意味著它不太可能與靜態分析工具良好協作,這些工具可能會報告無效的參數。同樣,可讀性方面也付出了代價:這些可選參數不會包含在 NumPy 函數的 docstring 中。

目前尚不清楚這種權衡是否值得,因此我們建議暫時將其排除在外。添加實作特定的參數將需要直接使用這些函式庫。

協定的其他可能選擇#

陣列函數 __array_function__ 僅包含兩個參數,functypes,它們提供有關函數呼叫上下文的資訊。

func 是協定的一部分,因為沒有辦法避免它:實作需要能夠透過將函數與 NumPy 的公共 API 匹配來進行分派。

包含 types 是因為我們可以幾乎免費地計算它,作為收集要在 implement_array_function 中呼叫的 __array_function__ 實作的一部分。我們也認為許多 __array_function__ 方法會使用它,否則它們將需要自己提取此資訊。提供每種類型的單個實例同樣容易,但僅提供類型似乎更清晰。

更進一步來說,有人建議 __array_function__ 應該是一個 classmethod。我們同意刪除冗餘的 self 參數會更簡潔一些,但覺得這種微小的清理不值得打破 __array_ufunc__ 的先例。

我們認為可能需要傳遞給 __array_ufunc__ 實作的其他兩個參數是

  • 存取非分派實作 (即,在用 array_function_dispatch 包裝之前) 在 ndarray.__array_function__ 中,將允許我們從 implement_array_function 中刪除該方法的特殊情況邏輯。

  • 存取傳遞到 array_function_dispatch() 中的 dispatcher 函數,將允許 __array_function__ 實作透過呼叫 dispatcher(*args, **kwargs) 以通用方式確定「類陣列」參數的列表。這對於基於陣列屬性 (例如,dtypeunits) 的值而不是直接基於陣列類型進行分派的 __array_function__ 實作可能很有用。

我們目前已將這些排除在外,因為我們不知道它們是必要的。如果我們想在未來包含它們,最簡單的方法是更新 array_function_dispatch 裝飾器,將它們添加為函數屬性。

在執行時產生的可呼叫物件#

NumPy 有一些 API 以動態方式定義可呼叫物件,例如 vectorizerandom.RandomState 物件上的方法。在科學 Python 堆疊中的其他核心函式庫中也可以找到範例,例如 scipy.stats 中的分佈物件和 scikit-learn 中的模型物件。如果也能夠為此類可呼叫物件編寫重載,那就太好了。這對 __array_function__ 協定提出了挑戰,因為與函數的情況不同,在 numpy 命名空間中沒有公共物件可以傳遞到 func 參數中。

我們可以透過為如何檢查 func 參數建立替代約定來潛在地處理這個問題,例如,透過使用 func.__self__ 來獲取類別物件,並使用 func.__func__ 返回未綁定的函數物件。但是,需要謹慎,因為這會將目前作為介面永久功能的實作細節,例如 vectorize 是作為類別而不是閉包實作的事實,或者方法是直接實作還是使用描述器實作的事實,都混雜在一起。

鑑於複雜性和有限的用例,我們也暫時擱置這個問題,但我們相信,如果需要,__array_function__ 可以擴展以適應這些用例。

討論#

此提案的各種替代方案在幾個 GitHub issue 中進行了討論

  1. pydata/sparse #1

  2. numpy/numpy #11129

此外,它還是 一篇部落格文章的主題。在此之後,它在 NumPy 開發者衝刺會議 上進行了討論,該會議在 加州大學柏克萊分校資料科學研究所 (BIDS) 舉行。

有關此提案本身的詳細討論,可以在 郵件列表 和相關的提取請求 (1, 2, 3) 中找到