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 de 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 a la misma clasificación que la matriz X, replicando 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 replicando en las 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 emisión.
Principios
El lenguaje 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 preparadas 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 en las operaciones entre arrays de diferentes clasificaciones. Esto es diferente de NumPy, que infiere la especificación cuando es posible.
Cómo transmitir un array de rango inferior a uno de rango más alto
Los escalares siempre se pueden transmitir en arreglos sin una especificación explícita de dimensiones de transmisión. Una operación binaria a nivel de elementos entre un escalar y un array significa aplicar la operación con el escalar a cada elemento del array. 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 transmisió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 las dimensiones del arreglo de rango más alto que 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, significa que 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 x 3 (dimensiones (3,3)). Hay dos formas de realizar la transmisión 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 de 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 describa cómo una forma de rango más pequeña se transmite en una forma de rango más grande. Por ejemplo, en un cuboide de 2×3×4 y una matriz de 3×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 array de rango inferior con uno de rango superior especificando con qué dimensiones del array de rango superior 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 arreglo 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 de la tupla de transmisión debe ser el mismo en el que se espera que las dimensiones del arreglo de rango inferior coincidan con las dimensiones del arreglo de rango más alto. El primer elemento de la tupla especifica qué dimensión del arreglo de rango superior debe coincidir con la dimensión 0 en el arreglo de rango más bajo. El segundo elemento de la tupla especifica qué dimensión del arreglo de rango superior 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 arreglos 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) son incompatibles y no pueden transmitirse.
Surge un caso especial (y también se admite) 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) transmisión 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 bajo a uno de rango más alto y la transmisión con dimensiones degeneradas 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 juntas 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) con 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 de 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, "degeneración de la 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 de 1 x 2 que se agrega a un array de tamaño de 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 emisión para producir un array 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 de 4x3x1), lo que produce un array intermedio de 4x1x2. 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 1 x 2 como lo son las dimensiones de transmisión (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 arreglo de 4 × 3 × 2.