本頁面由 Cloud Translation API 翻譯而成。
Switch to English

廣播語義

本文檔描述了XLA中的廣播語義如何工作。

什麼是廣播?

廣播是使具有不同形狀的數組具有兼容形狀以進行算術運算的過程。該術語是從Numpy 廣播借來的。

在不同等級的多維數組之間或具有不同但兼容形狀的多維數組之間的操作可能需要廣播。考慮加法X+v ,其中X是矩陣(秩2的數組),而v是向量(秩1的數組)。為了執行逐元素加法,XLA需要通過複製v一定次數來將向量v “廣播”到與矩陣X相同的等級。向量的長度必須至少與矩陣的尺寸之一匹配。

例如:

 |1 2 3| + |7 8 9|
|4 5 6|
 

矩陣的維數是(2,3),向量的維數是(3)。通過在行上複製向量來廣播向量,從而得到:

 |1 2 3| + |7 8 9| = |8  10 12|
|4 5 6|   |7 8 9|   |11 13 15|
 

在Numpy中,這稱為廣播

原則

XLA語言盡可能嚴格和顯式,避免了隱式和“魔術”功能。這樣的功能可能會使一些計算的定義更容易些,但代價是要在用戶代碼中增加更多的假設,而這些假設在長期內將很難更改。如有必要,可以在客戶端級別的包裝程序中添加隱式和魔術功能。

關於廣播,需要關於不同等級的陣列之間的操作的明確的廣播規範。這與Numpy不同,後者會在可能的情況下推斷出規格。

將較低級別的數組廣播到較高級別的數組

標量可以始終在數組中廣播,而無需明確說明廣播維度。標量和數組之間的逐元素二進制運算意味著對數組中的每個元素應用標量運算。例如,將標量添加到矩陣意味著生成一個矩陣,該矩陣的每個元素都是標量與相應輸入矩陣元素的和。

 |1 2 3| + 7 = |8  9  10|
|4 5 6|       |11 12 13|
 

可以通過在二進制運算中使用元組來捕獲大多數廣播需求。當操作的輸入具有不同的等級時,此廣播元組指定較高等級數組中的哪個維度與較低等級數組匹配。

考慮前面的示例,而不是將標量添加到(2,3)矩陣,而是將維數(3)的向量添加到維數(2,3)的矩陣。 如果不指定廣播,則此操作無效。為了正確地請求矩陣向量相加,請將廣播維度指定為(1),這意味著向量的維度與矩陣的維度1相匹配。在2D中,如果將維度0視為行,將維度1視為列,則意味著向量的每個元素將成為一列,其大小與矩陣中的行數匹配:

 |7 8 9| ==> |7 8 9|
            |7 8 9|
 

作為更複雜的示例,請考慮將3元素向量(維度(3))添加到3x3矩陣(維度(3,3))。對於此示例,可以通過兩種方式進行廣播:

(1)可以使用1的廣播尺寸。每個向量元素變為一列,並且向量在矩陣中的每一行都重複。

 |7 8 9| ==> |7 8 9|
            |7 8 9|
            |7 8 9|
 

(2)可以使用廣播維度0。每個向量元素變為一行,並且該向量對於矩陣中的每一列都是重複的。

  |7| ==> |7 7 7|
 |8|     |8 8 8|
 |9|     |9 9 9|
 

廣播維度可以是元組,其描述瞭如何將較小秩的形狀廣播為較大秩的形狀。例如,給定2x3x4長方體和3x4矩陣,廣播元組(1,2)意味著將矩陣與長方體的尺寸1和2匹配。

如果指定了broadcast_dimensions參數,則在XlaBuilder的二進制操作中使用這種廣播類型。例如,請參見XlaBuilder :: Add 。在XLA源代碼中,這種廣播有時稱為“ InDim”廣播。

正式定義

廣播屬性允許通過指定要匹配的較高維度數組的維度來匹配較低排名數組與較高排名數組。例如,對於維數為MxNxPxQ的數組,維數為T的向量可以按如下方式進行匹配:

           MxNxPxQ

dim 3:          T
dim 2:        T
dim 1:      T
dim 0:    T
 

在每種情況下,T必須等於較高級別數組的匹配維度。然後,向量值將從匹配的維度廣播到所有其他維度。

要將TxV矩陣匹配到MxNxPxQ陣列,使用了一對廣播尺寸:

           MxNxPxQ
dim 2,3:      T V
dim 1,2:    T V
dim 0,3:  T     V
etc...
 

廣播元組中的尺寸順序必須是期望下排陣列的尺寸與上排陣列的尺寸匹配的順序。元組中的第一個元素表示較高級別數組中的哪個維度必須與較低級別數組中的維度0相匹配。維度1的第二個元素,依此類推。廣播維度的順序必須嚴格增加。例如,在前面的示例中,將V匹配為N,將T匹配為P是非法的;同時將V與P和N匹配也是非法的。

廣播尺寸退化的相似秩的數組

一個相關的廣播問題是廣播兩個具有相同等級但尺寸大小不同的陣列。與Numpy的規則類似,這僅在數組兼容時才可行。當兩個數組的所有尺寸都兼容時,它們就是兼容的。在以下情況下,兩個維度兼容:

  • 它們相等,或者
  • 其中之一是1(“退化”維)

當遇到兩個兼容的數組時,結果形狀在每個維度索引的兩個輸入中具有最大值。

例子:

  1. (2,1)和(2,3)廣播到(2,3)。
  2. (1,2,5)和(7,2,5)廣播到(7,2,5)
  3. (7,2,5)和(7,1,5)廣播到(7,2,5)
  4. (7,2,5)和(7,2,6)不兼容,無法廣播。

出現並支持一種特殊情況,其中每個輸入數組在不同的索引處具有退化的維數。在這種情況下,結果是“外部操作”:(2,1)和(1,3)廣播到(2,3)。有關更多示例,請參閱有關廣播Numpy文檔

廣播組成

可以在相同的二進制操作中執行將較低等級的陣列廣播到較高等級的陣列以及使用退化維進行廣播。例如,可以使用廣播維度值(0)將大小為4的向量和大小為1x2的矩陣相加:

 |1 2 3 4| + [5 6]    // [5 6] is a 1x2 matrix, not a vector.
 

首先,使用廣播維度將向量廣播到等級2(矩陣)。廣播維中的單個值(0)表示矢量的維零與矩陣的維零匹配。這將生成一個大小為4xM的矩陣,其中選擇值M以匹配1x2數組中相應的尺寸大小。因此,生成了一個4x2矩陣:

 |1 1| + [5 6]
|2 2|
|3 3|
|4 4|
 

然後,“退化維廣播”廣播1x2矩陣的維零,以匹配右側的相應維大小:

 |1 1| + |5 6|     |6  7|
|2 2| + |5 6|  =  |7  8|
|3 3| + |5 6|     |8  9|
|4 4| + |5 6|     |9 10|
 

一個更複雜的示例是使用廣播維數(1、2)將大小為1x2的矩陣添加到大小為4x3x1的數組中。首先,使用廣播維度將1x2矩陣廣播到第3級,以產生中間Mx1x2陣列,其中維度大小M由產生4x1x2中間陣列的較大操作數(4x3x1陣列)的大小確定。 M的尺寸為0(最左側的尺寸),因為當廣播尺寸為(1、2)時,尺寸1和2映射到原始1x2矩陣的尺寸。可以使用退化尺寸的廣播將此中間陣列添加到4x3x1矩陣中,以產生4x3x2陣列結果。