Вещание

В этом документе описывается семантика вещания XLA.

Что такое вещание?

Рассылка — это процесс придания массивам различной формы совместимых форм для арифметических операций. Терминология заимствована из трансляции NumPy .

Широковещательная передача может потребоваться для операций между многомерными массивами разных рангов или между многомерными массивами разных, но совместимых форм. Рассмотрим сложение X+v где X — матрица (двухмерный массив), а v — вектор (одномерный массив). Чтобы выполнить поэлементное сложение, XLA необходимо «транслировать» вектор v в то же количество измерений, что и матрица X , путем репликации v определенное количество раз. Длина вектора должна совпадать хотя бы с одним из размеров матрицы.

Например:

|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|

Большинство потребностей в вещании можно удовлетворить, используя кортеж измерений в двоичной операции. Когда входные данные операции имеют разные ранги, этот широковещательный кортеж определяет, какие измерения в массиве более высокой размерности должны соответствовать массиву меньшей размерности .

Рассмотрим предыдущий пример. Вместо добавления скаляра к матрице (2,3) добавьте вектор размерности (3) к матрице размерностей (2,3). Без указания широковещания эта операция недействительна. Чтобы правильно запросить сложение матрицы и вектора, укажите размерность широковещательной передачи (1), что означает, что размерность вектора соответствует размерности 1 матрицы. В 2D, если измерение 0 представляет строки, а измерение 1 представляет столбцы, это означает, что каждый элемент вектора становится столбцом, размер которого соответствует количеству строк в матрице:

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

В качестве более сложного примера рассмотрим добавление трехэлементного вектора (размерность (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 кубоида.

Этот тип трансляции используется в двоичных операциях в 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 по трансляции .

Состав трансляции

Трансляция массива меньшей размерности в массив более высокой размерности и трансляция с использованием вырожденных измерений могут выполняться в одной и той же двоичной операции. Например, вектор размером 4 и матрицу размером 1x2 можно сложить вместе, используя широковещательные измерения со значением (0):

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

Сначала вектор транслируется до двумерного (матричного) размера с использованием широковещательных измерений. Единственное значение (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|

Более сложный пример — матрица размера 1x2, добавленная к массиву размером 4x3x1 с использованием широковещательных размеров (1, 2). Сначала матрица 1x2 транслируется до трех измерений с использованием измерений широковещания для создания промежуточного массива Mx1x2, где размер измерения M определяется размером большего операнда (массива 4x3x1), создающего промежуточный массив 4x1x2. M находится в измерении 0 (самое левое измерение), поскольку измерения 1 и 2 сопоставляются с размерами исходной матрицы 1x2, поскольку размеры широковещательной передачи равны (1, 2). Этот промежуточный массив можно добавить к матрице 4x3x1, используя трансляцию вырожденных измерений, чтобы получить результат массива 4x3x2.