Este documento descreve a semântica de transmissão do XLA.
O que é transmissão?
A transmissão é o processo de fazer com que matrizes com formas diferentes tenham formas compatíveis com operações aritméticas. A terminologia é emprestada da transmissão numpy.
A transmissão pode ser necessária para operações entre matrizes multidimensionais de classificações diferentes ou entre matrizes multidimensionais com formas diferentes, mas compatíveis. Considere a adição X+v, em que X é uma matriz (uma matriz de classificação 2) e v é um vetor (uma matriz de classificação 1). Para realizar a adição por elemento, o XLA precisa "transmitir" o vetor v para a mesma classificação que a matriz X, replicando v um determinado número de vezes. O comprimento do vetor
precisa corresponder a pelo menos uma das dimensões da matriz.
Exemplo:
|1 2 3| + |7 8 9|
|4 5 6|
As dimensões da matriz são (2,3) e a dimensão do vetor é (3). O vetor é transmitido replicando-o em linhas para conseguir:
|1 2 3| + |7 8 9| = |8 10 12|
|4 5 6| |7 8 9| |11 13 15|
No NumPy, isso é chamado de transmissão.
Princípios
A linguagem XLA é a mais rigorosa e explícita possível, evitando recursos "mágicos" implícitos. Esses recursos podem facilitar a definição de alguns cálculos, mas à custa de mais suposições incorporadas ao código do usuário que serão difíceis de alterar a longo prazo. Se necessário, é possível adicionar recursos mágicos implícitos a wrappers no nível do cliente.
Com relação à transmissão, o XLA exige especificações explícitas de transmissão em operações entre matrizes de classificações diferentes. Ele é diferente de NumPy, que infere a especificação quando possível.
Transmitir uma matriz de classificação inferior em uma matriz de classificação superior
Os escalars sempre podem ser transmitidos em matrizes sem uma especificação explícita das dimensões de transmissão. Uma operação binária de elemento entre um escalar e uma matriz significa aplicar a operação com o escalar a cada elemento da matriz. Por exemplo, adicionar um escalar a uma matriz significa produzir uma matriz em que cada elemento é uma soma do escalar e do elemento correspondente da matriz de entrada.
|1 2 3| + 7 = |8 9 10|
|4 5 6| |11 12 13|
A maioria das necessidades de transmissão pode ser capturada usando uma tupla de dimensões em uma operação binária. Quando as entradas da operação têm classificações diferentes, essa tupla de transmissão especifica quais dimensões na matriz high-rank precisam corresponder à matriz lower-rank.
Considere o exemplo anterior. Em vez de adicionar uma matriz escalar a uma (2,3), adicione um vetor de dimensão (3) a uma matriz de dimensões (2,3). Sem especificar a transmissão, essa operação é inválida. Para solicitar corretamente a adição de vetor de matriz, especifique a dimensão de transmissão como (1), o que significa que a dimensão do vetor corresponde à dimensão 1 da matriz. No 2D, se a dimensão 0 representa linhas e a dimensão 1 representa colunas, isso significa que cada elemento do vetor se torna uma coluna de tamanho correspondente ao número de linhas na matriz:
|7 8 9| ==> |7 8 9|
|7 8 9|
Como um exemplo mais complexo, adicione um vetor de três elementos (dimensão (3)) a uma matriz 3x3 (dimensões (3,3)). Há duas maneiras de a transmissão acontecer para este exemplo:
(1) Uma dimensão de transmissão 1 pode ser usada. Cada elemento vetorial se torna uma coluna, e o vetor é duplicado em cada linha na matriz.
|7 8 9| ==> |7 8 9|
|7 8 9|
|7 8 9|
(2) Uma dimensão de transmissão de 0 pode ser usada. Cada elemento vetorial se torna uma linha, e o vetor é duplicado em cada coluna na matriz.
|7| ==> |7 7 7|
|8| |8 8 8|
|9| |9 9 9|
As dimensões de transmissão podem ser uma tupla que descreve como uma forma de classificação menor é transmitida para um formato de classificação maior. Por exemplo, considerando um cuboide de 2 x 3 x 4 e uma matriz 3 x 4, uma tupla de transmissão (1,2) significa fazer a correspondência da matriz com as dimensões 1 e 2 do cubo.
Esse tipo de transmissão é usado nas operações binárias em XlaBuilder quando o
argumento broadcast_dimensions é fornecido. Por exemplo, consulte
XlaBuilder::Add.
No código-fonte do XLA, esse tipo de transmissão às vezes é chamado de transmissão "InDim".
Definição formal
O atributo de transmissão permite corresponder uma matriz de classificação inferior a uma de classificação superior especificando quais dimensões da matriz de classificação mais alta corresponder. Por exemplo, para uma matriz com dimensões MxNxPxQ, um vetor com a dimensão T pode ser correspondido da seguinte maneira:
MxNxPxQ
dim 3: T
dim 2: T
dim 1: T
dim 0: T
Em cada caso, T precisa ser igual à dimensão correspondente da matriz de classificação mais alta. Os valores do vetor são transmitidos da dimensão correspondente para todas as outras dimensões.
Para fazer a correspondência de uma matriz TxV com a matriz MxNxPxQ, é usado um par de dimensões de transmissão:
MxNxPxQ
dim 2,3: T V
dim 1,2: T V
dim 0,3: T V
etc...
A ordem das dimensões na tupla de transmissão precisa ser a mesma em que as dimensões da matriz de classificação mais baixa devem corresponder às dimensões da matriz de classificação mais alta. O primeiro elemento na tupla especifica qual dimensão na matriz de classificação mais alta deve corresponder à dimensão 0 na matriz de classificação inferior. O segundo elemento na tupla especifica qual dimensão na matriz de classificação mais alta deve corresponder à dimensão 1 na matriz de classificação inferior e assim por diante. A ordem das dimensões de transmissão precisa ser estritamente crescente. Por exemplo, no exemplo anterior, é ilegal combinar V para N e T para P; também é ilegal combinar V a P e N.
Transmissão de matrizes de classificação semelhante com dimensões desgeneradas
Um problema relacionado é transmitir duas matrizes com a mesma classificação, mas tamanhos de dimensão diferentes. Assim como o NumPy, isso só é possível quando as matrizes são compatíveis. Duas matrizes são compatíveis quando todas as dimensões delas são compatíveis. Duas dimensões são compatíveis se:
- Eles são iguais ou
- Um deles é 1 (uma dimensão "degenerar")
Quando duas matrizes compatíveis são encontradas, a forma de resultado tem o máximo das duas entradas em cada índice de dimensão.
Exemplos:
- (2,1) e (2,3) são transmitidos para (2,3).
- (1,2,5) e (7,2,5) são transmitidas para (7,2,5).
- (7,2,5) e (7,1,5) são transmitidas para (7,2,5).
- (7,2,5) e (7,2,6) são incompatíveis e não podem ser transmitidos.
Há um caso especial que também é aceito, em que cada uma das matrizes de entrada tem uma dimensão "degenerada" em um índice diferente. Nesse caso, o resultado é uma "operação externa": (2,1) e (1,3) é transmitida para (2,3). Para mais exemplos, consulte a documentação do NumPy sobre transmissão.
Composição da transmissão
A transmissão de uma matriz de classificação inferior para uma matriz de classificação mais alta e a transmissão usando dimensões desgeneradas podem ser realizadas na mesma operação binária. Por exemplo, um vetor de tamanho 4 e uma matriz de tamanho 1x2 podem ser somados usando dimensões de transmissão de valor (0):
|1 2 3 4| + [5 6] // [5 6] is a 1x2 matrix, not a vector.
Primeiro, o vetor é transmitido até a classificação 2 (matriz) usando as dimensões de transmissão. O único valor (0) nas dimensões de transmissão indica que a dimensão zero do vetor corresponde à dimensão zero da matriz. Isso produz uma matriz de tamanho 4 x M, em que o valor M é escolhido para corresponder ao tamanho da dimensão correspondente na matriz de 1 x 2. Portanto, é produzida uma matriz 4x2:
|1 1| + [5 6]
|2 2|
|3 3|
|4 4|
Em seguida, "desgerar transmissão de dimensão" transmite a dimensão zero da matriz 1x2 para corresponder ao tamanho de dimensão correspondente do lado direito:
|1 1| + |5 6| |6 7|
|2 2| + |5 6| = |7 8|
|3 3| + |5 6| |8 9|
|4 4| + |5 6| |9 10|
Um exemplo mais complicado é uma matriz de tamanho 1 x 2 adicionada a uma matriz de tamanho 4 x 3 x 1 usando dimensões de transmissão de (1, 2). Primeiro, a matriz 1x2 é transmitida até a classificação 3 usando as dimensões de transmissão para produzir uma matriz Mx1x2 intermediária, em que o tamanho da dimensão M é determinado pelo tamanho do operando maior (a matriz 4x3x1), produzindo uma matriz intermediária 4x1x2. O M está na dimensão 0 (a mais à esquerda) porque as dimensões 1 e 2 estão mapeadas para as dimensões da matriz 1x2 original, assim como as dimensões de transmissão (1, 2). Essa matriz intermediária pode ser adicionada à matriz 4x3x1 usando a transmissão de dimensões desgeneradas para produzir um resultado de matriz de 4x3x2.