En este documento, se describe la semántica de transmisión de XLA.
¿Qué es una transmisión?
La transmisión es el proceso de hacer que los arrays con diferentes formas tengan formas compatibles para las operaciones aritméticas. La terminología se toma de la transmisión NumPy.
Es posible que se requiera la transmisión para operaciones entre arrays multidimensionales de diferentes rangos o entre arrays multidimensionales con formas diferentes pero compatibles. Considera la suma X+v, en la que X es una matriz (un arreglo de rango 2) y v es un vector (un arreglo de rango 1). Para realizar una adición a nivel de los elementos, XLA debe "transmitir" el vector v al mismo rango que la matriz X, mediante la replicación de v una cierta cantidad de veces. La longitud del vector debe coincidir con al menos una de las dimensiones de la matriz.
Por ejemplo:
|1 2 3| + |7 8 9|
|4 5 6|
Las dimensiones de la matriz son (2,3) y la dimensión del vector es (3). El vector se transmite por medio de su replicación en filas para obtener lo siguiente:
|1 2 3| + |7 8 9| = |8 10 12|
|4 5 6| |7 8 9| |11 13 15|
En NumPy, esto se llama transmisión.
Principios
El lenguaje de XLA es lo más estricto y explícito posible, lo que evita las funciones "mágicas" implícitas. Estas funciones pueden hacer que algunos cálculos sean un poco más fáciles de definir, pero a costa de más suposiciones incorporadas en el código de usuario que serán difíciles de cambiar a largo plazo. Si es necesario, se pueden agregar funciones mágicas implícitas en los wrappers a nivel del cliente.
Con respecto a la transmisión, XLA requiere especificaciones de transmisión explícitas sobre las operaciones entre arrays de diferentes clasificaciones. Esto es diferente de NumPy, que infiere la especificación cuando es posible.
Transmite un array de menor rango a uno de mayor rango
Los escalares siempre se pueden transmitir a través de arrays sin una especificación explícita de dimensiones de transmisión. Una operación binaria a nivel de elementos entre un escalar y un arreglo significa aplicar la operación con el escalar a cada elemento del arreglo. Por ejemplo, agregar un escalar a una matriz significa producir una matriz en la que cada elemento es una suma del escalar y el elemento correspondiente de la matriz de entrada.
|1 2 3| + 7 = |8 9 10|
|4 5 6| |11 12 13|
La mayoría de las necesidades de emisión se pueden capturar con una tupla de dimensiones en una operación binaria. Cuando las entradas a la operación tienen diferentes clasificaciones, esta tupla de transmisión especifica qué dimensiones del arreglo de rango superior deben coincidir con el arreglo de rango más bajo.
Considera el ejemplo anterior. En lugar de agregar un escalar a una matriz (2,3), agrega un vector de dimensión (3) a una matriz de dimensiones (2,3). Si no se especifica la transmisión, esta operación no es válida. Para solicitar correctamente la adición del vector de matriz, especifica que la dimensión de transmisión sea (1), lo que significa que la dimensión del vector coincide con la dimensión 1 de la matriz. En 2D, si la dimensión 0 representa filas y la dimensión 1 representa columnas, cada elemento del vector se convierte en una columna de un tamaño que coincide con el número de filas en la matriz:
|7 8 9| ==> |7 8 9|
|7 8 9|
Como ejemplo más complejo, considera agregar un vector de 3 elementos (dimensión (3)) a una matriz de 3 × 3 (dimensiones (3,3)). Existen dos formas de transmitir en este ejemplo:
(1) Se puede usar una dimensión de transmisión igual a 1. Cada elemento vectorial se convierte en una columna y el vector se duplica para cada fila de la matriz.
|7 8 9| ==> |7 8 9|
|7 8 9|
|7 8 9|
(2) Se puede usar una dimensión de transmisión igual a 0. Cada elemento vectorial se convierte en una fila y el vector se duplica para cada columna de la matriz.
|7| ==> |7 7 7|
|8| |8 8 8|
|9| |9 9 9|
Las dimensiones de transmisión pueden ser una tupla que describe cómo una forma de rango más pequeño se transmite a una forma de rango más grande. Por ejemplo, con un cuboide de 2 x 3 x 4 y una matriz de 3 x 4, una tupla de transmisión (1,2) significa hacer coincidir la matriz con las dimensiones 1 y 2 del cuboide.
Este tipo de transmisión se usa en las operaciones binarias de XlaBuilder si se proporciona el argumento broadcast_dimensions. Por ejemplo, consulta XlaBuilder::Add.
En el código fuente de XLA, este tipo de transmisión a veces se denomina transmisión "InDim".
Definición formal
El atributo de transmisión permite hacer coincidir un arreglo de rango inferior con uno de rango superior especificando las dimensiones del arreglo de rango superior con las que coincidir. Por ejemplo, para un array con dimensiones MxNxPxQ, un vector con dimensión T puede coincidir de la siguiente manera:
MxNxPxQ
dim 3: T
dim 2: T
dim 1: T
dim 0: T
En cada caso, T debe ser igual a la dimensión coincidente del array de rango superior. Luego, los valores del vector se transmiten desde la dimensión coincidente a todas las demás dimensiones.
Para hacer coincidir una matriz TxV con el array MxNxPxQ, se usa un par de dimensiones de transmisión:
MxNxPxQ
dim 2,3: T V
dim 1,2: T V
dim 0,3: T V
etc...
El orden de las dimensiones en la tupla de emisión debe ser el mismo en el que se espera que las dimensiones del array de rango inferior coincidan con las dimensiones del array de rango más alto. El primer elemento de la tupla especifica qué dimensión del array de rango más alto debe coincidir con la dimensión 0 en el array de rango más bajo. El segundo elemento de la tupla especifica qué dimensión del arreglo de rango más alto debe coincidir con la dimensión 1 en el arreglo de rango inferior, y así sucesivamente. El orden de las dimensiones de transmisión debe aumentar estrictamente. Por ejemplo, en el ejemplo anterior es ilegal hacer coincidir V con N y T con P; también es ilegal hacer coincidir V con P y N.
Cómo transmitir arrays de rango similar con dimensiones degeneradas
Un problema relacionado es la transmisión de dos arrays que tienen la misma clasificación, pero diferentes tamaños de dimensión. Al igual que con NumPy, esto solo es posible cuando los arreglos son compatibles. Dos arrays son compatibles cuando todas sus dimensiones también lo son. Dos dimensiones son compatibles en los siguientes casos:
- Son iguales o
- Uno de ellos es 1 (una dimensión de "degeneración").
Cuando se encuentran dos arrays compatibles, la forma del resultado tiene el máximo de las dos entradas en cada índice de dimensión.
Ejemplos:
- (2,1) y (2,3) transmiten a (2,3).
- (1,2,5) y (7,2,5) transmiten a (7,2,5).
- (7,2,5) y (7,1,5) transmiten a (7,2,5).
- (7,2,5) y (7,2,6) no son compatibles y no se pueden transmitir.
Surge un caso especial, y también es compatible, en el que cada uno de los arreglos de entrada tiene una dimensión en un índice diferente. En este caso, el resultado es una "operación externa": (2,1) y (1,3) se transmite a (2,3). Para obtener más ejemplos, consulta la documentación de NumPy sobre transmisiones.
Composición de la transmisión
La transmisión de un array de rango inferior a un array de rango superior y la transmisión mediante dimensiones de degeneración se pueden realizar en la misma operación binaria. Por ejemplo, un vector de tamaño 4 y una matriz de tamaño 1 x 2 se pueden sumar usando dimensiones de transmisión de valor (0):
|1 2 3 4| + [5 6] // [5 6] is a 1x2 matrix, not a vector.
Primero, el vector se transmite hasta el rango 2 (matriz) mediante las dimensiones de transmisión. El valor único (0) en las dimensiones de transmisión indica que la dimensión cero del vector coincide con la dimensión cero de la matriz. Esto produce una matriz de tamaño 4 x M en la que el valor M se elige para que coincida con el tamaño de dimensión correspondiente en el arreglo de 1 x 2. Por lo tanto, se produce una matriz de 4 x 2:
|1 1| + [5 6]
|2 2|
|3 3|
|4 4|
Luego, "degenerar transmisión de dimensiones" transmite la dimensión cero de la matriz de 1 x 2 para que coincida con el tamaño de dimensión correspondiente del lado derecho:
|1 1| + |5 6| |6 7|
|2 2| + |5 6| = |7 8|
|3 3| + |5 6| |8 9|
|4 4| + |5 6| |9 10|
Un ejemplo más complicado es una matriz de tamaño 1 x 2 agregada a un array de tamaño 4 x 3 x 1 con dimensiones de transmisión de (1, 2). Primero, la matriz 1x2 se transmite hasta el rango 3 con las dimensiones de transmisión para producir un arreglo Mx1x2 intermedio, en el que el tamaño de dimensión M se determina por el tamaño del operando más grande (el array 4x3x1) produciendo un array intermedio de 4x1x2. La M está en la dimensión 0 (la dimensión más a la izquierda) porque las dimensiones 1 y 2 se asignan a las dimensiones de la matriz original de 1x2 como las dimensiones de transmisión son (1, 2). Este array intermedio se puede agregar a la matriz de 4 × 3 × 1 mediante la transmisión de dimensiones degeneradas para producir un resultado de array de 4 × 3 × 2.