Семантика вещания

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

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

Широковещательная рассылка - это процесс создания массивов различной формы с совместимыми формами для арифметических операций. Терминология заимствована из вещания Numpy.

Широковещательная передача может потребоваться для операций между многомерными массивами разных рангов или между многомерными массивами с разными, но совместимыми формами. Рассмотрим сложение X+v где X - матрица (массив ранга 2), а v - вектор (массив ранга 1). Чтобы выполнить поэлементное сложение, 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 является максимально строгим и явным, избегая неявных и «волшебных» функций. Такие функции могут немного облегчить определение некоторых вычислений за счет большего количества предположений, встроенных в код пользователя, которые будет трудно изменить в долгосрочной перспективе. При необходимости в оболочки клиентского уровня можно добавить неявные и волшебные функции.

Что касается широковещательной передачи, требуются явные спецификации широковещательной передачи для операций между массивами разных рангов. Это отличается от 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 кубоида.

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

Сначала вектор транслируется до ранга 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|

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