В этом документе описывается, как работает семантика вещания в 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 , если 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.