廣播

本文件說明 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 語言盡可能嚴格且明確,避免使用隱含的「神秘」功能。這類功能可能會讓某些運算較容易定義,但實際上,使用者程式碼會編寫較多假設,而長期來說難以改變。如有需要,您可以在用戶端層級的包裝函式中新增隱含的神奇功能。

針對廣播,XLA 需要針對不同排名陣列之間的作業提供明確播送規格。這與 NumPy 不同,後者會盡可能推斷規格。

將較低排名陣列播送到較高排名的陣列

「純量」一律可以透過陣列播送,而無須明確指定播送維度。純量和陣列之間的元素相關二進位運算代表將純量運算套用至陣列中的每個元素。例如,在矩陣中加入純量,意味著產生矩陣,每個元素都是純量和輸入矩陣對應元素的總和。

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

在二進位檔作業中使用多個維度,即可擷取大多數的播送需求。如果作業的輸入內容有不同的排名,此廣播元組會指定 higher-rank 陣列中的哪些維度要與 lower-rank 陣列比對。

請參考上一個範例。請勿將純量加入 (2,3) 矩陣,請在維度矩陣 (2,3) 中新增維度 (3) 。如果沒有指定廣播,這項作業就會無效。如要正確要求矩陣向量新增,請將播送維度指定為 (1),代表向量維度與矩陣的維度 1 相符。在 2D 中,如果維度 0 代表列,維度 1 代表欄,則代表向量的每個元素會變為大小與矩陣資料列數量相符的資料欄:

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

如果是更複雜的範例,請考慮在 3x3 矩陣 (維度 (3,3)) 中加入 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 的 cuboid 和 3x4 矩陣,廣播元組 (1,2) 代表將矩陣與 cuboid 的第 1 和 2 個維度進行比對。

這種廣播類型會用於 XlaBuilder 中的二進位作業 (如果指定了 broadcast_dimensions 引數)。例如,請參閱「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 陣列),由於廣播維度是 (1, 2),因此 M 位於維度 0 (最左側維度),因為維度 1 和 2 會對應至原始 1x2 矩陣的尺寸。此中繼陣列可以使用解壓縮尺寸播送,將 4x3x2 陣列加入 4x3x1 矩陣,進而產生 4x3x2 的陣列結果。