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