Semántica de operaciones

A continuación, se describe la semántica de las operaciones definidas en la interfaz de XlaBuilder. Por lo general, estas operaciones se asignan uno a uno a las operaciones definidas en la interfaz de RPC en xla_data.proto.

Nota sobre la nomenclatura: El tipo de datos generalizado con el que se encarga XLA es un array de N dimensiones que contiene elementos de algún tipo uniforme (como un número de punto flotante de 32 bits). En toda la documentación, array se usa para denotar un array de dimensiones arbitrarias. Para mayor comodidad, los casos especiales tienen nombres más específicos y conocidos. Por ejemplo, un vector es un array de 1 dimensión y una matriz es un array de 2 dimensiones.

AfterAll

Consulta también XlaBuilder::AfterAll.

AfterAll toma una cantidad variable de tokens y produce un solo token. Los tokens son tipos primitivos que se pueden agrupar en subprocesos entre operaciones de efectos secundarios para aplicar el orden. AfterAll se puede usar como una unión de tokens para pedir una operación después de una operación set.

AfterAll(operands)

Argumentos Tipo Semántica
operands XlaOp cantidad variable de tokens

AllGather

Consulta también XlaBuilder::AllGather.

Realiza la concatenación entre réplicas.

AllGather(operand, all_gather_dim, shard_count, replica_group_ids, channel_id)

Argumentos Tipo Semántica
operand XlaOp Es un array para concatenar entre réplicas.
all_gather_dim int64 Dimensión de concatenación
replica_groups vector de vectores de int64 Grupos entre los que se realiza la concatenación
channel_id int64 opcional ID de canal opcional para la comunicación entre módulos
  • replica_groups es una lista de grupos de réplicas entre los que se realiza la concatenación (el ID de réplica de la réplica actual se puede recuperar con ReplicaId). El orden de las réplicas en cada grupo determina el orden en el que se ubican sus entradas en el resultado. replica_groups debe estar vacío (en ese caso, todas las réplicas pertenecen a un solo grupo, ordenadas de 0 a N - 1) o contener la misma cantidad de elementos que la cantidad de réplicas. Por ejemplo, replica_groups = {0, 2}, {1, 3} realiza la concatenación entre las réplicas 0 y 2, y 1 y 3.
  • shard_count es el tamaño de cada grupo de réplicas. Esto lo necesitamos en los casos en que los objetos replica_groups están vacíos.
  • channel_id se usa para la comunicación entre módulos: solo las operaciones all-gather con el mismo channel_id pueden comunicarse entre sí.

La forma de la salida es la de entrada con el all_gather_dim hecho shard_count veces más grande. Por ejemplo, si hay dos réplicas y el operando tiene el valor [1.0, 2.5] y [3.0, 5.25] respectivamente en ambas réplicas, el valor de salida de esta op en la que all_gather_dim es 0 será [1.0, 2.5, 3.0, 5.25] en ambas réplicas.

AllReduce

Consulta también XlaBuilder::AllReduce.

Realiza un procesamiento personalizado en las réplicas.

AllReduce(operand, computation, replica_group_ids, channel_id)

Argumentos Tipo Semántica
operand XlaOp un array o una tupla no vacía de arrays para reducir en las réplicas
computation XlaComputation Cálculo de reducción
replica_groups vector de vectores de int64 Grupos entre los que se realizan las reducciones
channel_id int64 opcional ID de canal opcional para la comunicación entre módulos
  • Cuando operand es una tupla de arreglos, se realiza el All-reduce en cada elemento de la tupla.
  • replica_groups es una lista de grupos de réplicas entre los que se realiza la reducción (el ID de réplica de la réplica actual se puede recuperar con ReplicaId). replica_groups debe estar vacío (en cuyo caso todas las réplicas pertenecen a un solo grupo) o contener la misma cantidad de elementos que la cantidad de réplicas. Por ejemplo, replica_groups = {0, 2}, {1, 3} realiza una reducción entre las réplicas 0 y 2, y 1 y 3.
  • channel_id se usa para la comunicación entre módulos: solo las operaciones all-reduce con el mismo channel_id pueden comunicarse entre sí.

La forma de salida es la misma que la de entrada. Por ejemplo, si hay dos réplicas y el operando tiene el valor [1.0, 2.5] y [3.0, 5.25], respectivamente, en las dos réplicas, el valor de salida de este cálculo de op y suma será [4.0, 7.75] en ambas réplicas. Si la entrada es una tupla, la salida también es una tupla.

Para calcular el resultado de AllReduce, es necesario tener una entrada de cada réplica, por lo que, si una réplica ejecuta un nodo AllReduce más veces que otra, la réplica anterior esperará por siempre. Dado que todas las réplicas ejecutan el mismo programa, no hay muchas formas de que eso suceda, pero es posible que la condición de un bucle while dependa de los datos de entrada, y los datos ingresados hagan que el bucle while se itere más veces en una réplica que en otra.

AllToAll

Consulta también XlaBuilder::AllToAll.

AllToAll es una operación colectiva que envía datos de todos los núcleos a todos los núcleos. Consta de dos fases:

  1. La fase de dispersión. En cada núcleo, el operando se divide en una cantidad de bloques split_count a lo largo del split_dimensions, y los bloques se dispersan en todos los núcleos, p.ej., el bloque i se envía al núcleo i.
  2. La fase de recopilación. Cada núcleo concatena los bloques recibidos a lo largo de concat_dimension.

Los núcleos que participan se pueden configurar de la siguiente manera:

  • replica_groups: Cada ReplicaGroup contiene una lista de los IDs de réplica que participan en el procesamiento (el ID de réplica de la réplica actual se puede recuperar con ReplicaId). AllToAll se aplicará dentro de los subgrupos en el orden especificado. Por ejemplo, replica_groups = { {1,2,3}, {4,5,0} } significa que se aplicará un AllToAll dentro de las réplicas {1, 2, 3} y en la fase de recopilación, y los bloques recibidos se concatenarán en el mismo orden de 1, 2, 3. Luego, se aplicará otro AllToAll en las réplicas 4, 5 y 0, y el orden de concatenación también es 4, 5, 0. Si replica_groups está vacío, todas las réplicas pertenecen a un grupo, en el orden de concatenación de su apariencia.

Requisitos previos:

  • El tamaño de la dimensión del operando en split_dimension es divisible por split_count.
  • La forma del operando no es una tupla.

AllToAll(operand, split_dimension, concat_dimension, split_count, replica_groups)

Argumentos Tipo Semántica
operand XlaOp array de entrada de dimensión n
split_dimension int64 Un valor en el intervalo [0, n) que nombra la dimensión en la que se divide el operando.
concat_dimension int64 Un valor en el intervalo [0, n) que nombra la dimensión a través de la cual se concatenan los bloques divididos
split_count int64 El número de núcleos que participan en esta operación. Si replica_groups está vacío, esta debería ser la cantidad de réplicas; de lo contrario, debería ser igual a la cantidad de réplicas en cada grupo.
replica_groups Vector de ReplicaGroup Cada grupo contiene una lista de IDs de réplica.

A continuación, se muestra un ejemplo de Alltoall.

XlaBuilder b("alltoall");
auto x = Parameter(&b, 0, ShapeUtil::MakeShape(F32, {4, 16}), "x");
AllToAll(x, /*split_dimension=*/1, /*concat_dimension=*/0, /*split_count=*/4);

En este ejemplo, hay 4 núcleos que participan en Alltoall. En cada núcleo, el operando se divide en 4 partes a lo largo de la dimensión 0, de modo que cada parte tiene la forma f32[4,4]. Las 4 partes están dispersas en todos los núcleos. Luego, cada núcleo concatena las partes recibidas en la dimensión 1, en el orden del núcleo del 0 al 4. Por lo tanto, el resultado en cada núcleo tiene la forma f32[16,4].

BatchNormGrad

Consulta también XlaBuilder::BatchNormGrad y el documento original de normalización por lotes para obtener una descripción detallada del algoritmo.

Calcula los gradientes de la norma por lotes.

BatchNormGrad(operand, scale, mean, variance, grad_output, epsilon, feature_index)

Argumentos Tipo Semántica
operand XlaOp matriz de dimensiones n que se normalizará (x).
scale XlaOp Array de 1 dimensión (\(\gamma\))
mean XlaOp Array de 1 dimensión (\(\mu\))
variance XlaOp Array de 1 dimensión (\(\sigma^2\))
grad_output XlaOp Gradientes pasados a BatchNormTraining (\(\nabla y\))
epsilon float Valor de Épsilon (\(\epsilon\))
feature_index int64 Índice a la dimensión del atributo en operand

Para cada atributo de la dimensión del atributo (feature_index es el índice de la dimensión del componente en operand), la operación calcula los gradientes con respecto a operand, offset y scale en todas las demás dimensiones. feature_index debe ser un índice válido para la dimensión del atributo en operand.

Los tres gradientes se definen mediante las siguientes fórmulas (si se supone que hay un array de 4 dimensiones como operand y que tiene el índice de dimensión de los atributos l, el tamaño del lote m y los tamaños espaciales w y h):

\[ \begin{split} c_l&= \frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h \left( \nabla y_{ijkl} \frac{x_{ijkl} - \mu_l}{\sigma^2_l+\epsilon} \right) \\\\ d_l&= \frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h \nabla y_{ijkl} \\\\ \nabla x_{ijkl} &= \frac{\gamma_{l} }{\sqrt{\sigma^2_{l}+\epsilon} } \left( \nabla y_{ijkl} - d_l - c_l (x_{ijkl} - \mu_{l}) \right) \\\\ \nabla \gamma_l &= \sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h \left( \nabla y_{ijkl} \frac{x_{ijkl} - \mu_l}{\sqrt{\sigma^2_{l}+\epsilon} } \right) \\\\\ \nabla \beta_l &= \sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h \nabla y_{ijkl} \end{split} \]

Las entradas mean y variance representan valores de momentos en dimensiones por lotes y espaciales.

El tipo de salida es una tupla de tres controladores:

Salidas Tipo Semántica
grad_operand XlaOp gradiente con respecto al operand de entrada ($\nabla x$)
grad_scale XlaOp gradiente con respecto al scale de entrada ($\nabla \gamma$)
grad_offset XlaOp gradiente con respecto a la entrada offset($\nabla \beta$)

BatchNormInference

Consulta también XlaBuilder::BatchNormInference y el documento original de normalización por lotes para obtener una descripción detallada del algoritmo.

Normaliza un array en dimensiones espaciales y por lotes.

BatchNormInference(operand, scale, offset, mean, variance, epsilon, feature_index)

Argumentos Tipo Semántica
operand XlaOp matriz de dimensiones n que se normalizará.
scale XlaOp Array de 1 dimensión
offset XlaOp Array de 1 dimensión
mean XlaOp Array de 1 dimensión
variance XlaOp Array de 1 dimensión
epsilon float Valor de Épsilon
feature_index int64 Índice a la dimensión del atributo en operand

Para cada atributo de la dimensión del componente (feature_index es el índice de la dimensión del componente en operand), la operación calcula la media y la varianza de todas las demás dimensiones y usa la media y la varianza para normalizar cada elemento en operand. feature_index debe ser un índice válido para la dimensión del componente en operand.

BatchNormInference equivale a llamar a BatchNormTraining sin procesar mean ni variance para cada lote. En su lugar, usa los valores mean y variance de entrada como valores estimados. El propósito de esta op es reducir la latencia en la inferencia, de ahí el nombre BatchNormInference.

El resultado es un array normalizado de n dimensiones con la misma forma que operand de entrada.

BatchNormTraining

Consulta también XlaBuilder::BatchNormTraining y the original batch normalization paper para obtener una descripción detallada del algoritmo.

Normaliza un array en dimensiones espaciales y por lotes.

BatchNormTraining(operand, scale, offset, epsilon, feature_index)

Argumentos Tipo Semántica
operand XlaOp matriz de dimensiones n que se normalizará (x).
scale XlaOp Array de 1 dimensión (\(\gamma\))
offset XlaOp Array de 1 dimensión (\(\beta\))
epsilon float Valor de Épsilon (\(\epsilon\))
feature_index int64 Índice a la dimensión del atributo en operand

Para cada atributo de la dimensión del componente (feature_index es el índice de la dimensión del componente en operand), la operación calcula la media y la varianza de todas las demás dimensiones y usa la media y la varianza para normalizar cada elemento en operand. feature_index debe ser un índice válido para la dimensión del componente en operand.

El algoritmo sigue el siguiente procedimiento para cada lote en operand \(x\) que contenga elementos m con w y h como el tamaño de las dimensiones espaciales (suponiendo que operand es un array de 4 dimensiones):

  • Calcula la media del lote \(\mu_l\) para cada atributo l en la dimensión del atributo: \(\mu_l=\frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h x_{ijkl}\)

  • Calcula la varianza del lote \(\sigma^2_l\): $\sigma^2l=\frac{1}{mwh}\sum{i=1}^m\sum{j=1}^w\sum{k=1}^h (x_{ijkl} - \mu_l)^2$

  • Normaliza, escala y cambia: \(y_{ijkl}=\frac{\gamma_l(x_{ijkl}-\mu_l)}{\sqrt[2]{\sigma^2_l+\epsilon} }+\beta_l\)

El valor épsilon, generalmente un número pequeño, se agrega para evitar errores de división por cero.

El tipo de salida es una tupla de tres XlaOp:

Salidas Tipo Semántica
output XlaOp array n dimensional con la misma forma que la entrada operand (y)
batch_mean XlaOp Array de 1 dimensión (\(\mu\))
batch_var XlaOp Array de 1 dimensión (\(\sigma^2\))

batch_mean y batch_var son momentos calculados a través de las dimensiones espaciales y del lote mediante las fórmulas anteriores.

BitcastConvertType

Consulta también XlaBuilder::BitcastConvertType.

Al igual que tf.bitcast en TensorFlow, realiza una operación de transmisión de bits a nivel de elementos desde una forma de datos hasta una forma objetivo. El tamaño de entrada y salida deben coincidir, p.ej., los elementos s32 se convierten en elementos f32 a través de la rutina de transmisión de bits, y un elemento s32 se convertirá en cuatro elementos s8. Bitcast se implementa como una conversión de bajo nivel, por lo que las máquinas con diferentes representaciones de punto flotante darán diferentes resultados.

BitcastConvertType(operand, new_element_type)

Argumentos Tipo Semántica
operand XlaOp matriz de tipo T con atenuación D
new_element_type PrimitiveType tipo U

Las dimensiones del operando y de la forma objetivo deben coincidir, además de la última dimensión, que cambiará según la proporción del tamaño primitivo antes y después de la conversión.

Los tipos de elementos de origen y destino no deben ser tuplas.

Conversión de bitcast a tipo primitivo de ancho diferente

La instrucción BitcastConvert de HLO admite el caso en el que el tamaño del tipo de elemento de salida T' no es igual al tamaño del elemento de entrada T. Dado que toda la operación es conceptualmente una conversión de bits y no cambia los bytes subyacentes, la forma del elemento de salida debe cambiar. Para B = sizeof(T), B' = sizeof(T'), hay dos casos posibles.

Primero, cuando es B > B', la forma de salida obtiene una nueva dimensión más pequeña de tamaño B/B'. Por ejemplo:

  f16[10,2]{1,0} %output = f16[10,2]{1,0} bitcast-convert(f32[10]{0} %input)

La regla sigue siendo la misma para los escalares eficaces:

  f16[2]{0} %output = f16[2]{0} bitcast-convert(f32[] %input)

Como alternativa, para B' > B, la instrucción requiere que la última dimensión lógica de la forma de entrada sea igual a B'/B, y esta dimensión se descarta durante la conversión:

  f32[10]{0} %output = f32[10]{0} bitcast-convert(f16[10,2]{1,0} %input)

Ten en cuenta que las conversiones entre diferentes anchos de bits no se realizan por elementos.

Transmisión

Consulta también XlaBuilder::Broadcast.

Agrega dimensiones a un array; para ello, duplica los datos en el array.

Broadcast(operand, broadcast_sizes)

Argumentos Tipo Semántica
operand XlaOp El array que se duplicará.
broadcast_sizes ArraySlice<int64> Los tamaños de las nuevas dimensiones

Las dimensiones nuevas se insertan a la izquierda; es decir, si broadcast_sizes tiene valores {a0, ..., aN} y la forma de operando tiene dimensiones {b0, ..., bM}, entonces la forma de la salida tiene las dimensiones {a0, ..., aN, b0, ..., bM}.

El nuevo índice de dimensiones en copias del operando; p.ej.,

output[i0, ..., iN, j0, ..., jM] = operand[j0, ..., jM]

Por ejemplo, si operand es un f32 escalar con el valor 2.0f y broadcast_sizes es {2, 3}, el resultado será un array con la forma f32[2, 3] y todos los valores en el resultado serán 2.0f.

BroadcastInDim

Consulta también XlaBuilder::BroadcastInDim.

Duplica el tamaño y la clasificación de un array; para ello, duplica los datos en él.

BroadcastInDim(operand, out_dim_size, broadcast_dimensions)

Argumentos Tipo Semántica
operand XlaOp El array que se duplicará.
out_dim_size ArraySlice<int64> Los tamaños de las dimensiones de la forma de destino
broadcast_dimensions ArraySlice<int64> ¿A qué dimensión de la forma objetivo corresponde cada dimensión de la forma de operando?

Es similar a Broadcast, pero permite agregar dimensiones en cualquier lugar y expandir las dimensiones existentes con tamaño 1.

operand se transmite a la forma que describe out_dim_size. broadcast_dimensions asigna las dimensiones de operand a las dimensiones de la forma objetivo, es decir, la dimensión i del operando se asigna a la dimensión broadcast_dimension[i] de la forma de salida. Las dimensiones de operand deben tener un tamaño de 1 o del mismo tamaño que la dimensión de la forma de salida a la que se asignan. Las dimensiones restantes se rellenan con dimensiones de tamaño 1. Luego, la transmisión de la dimensión de degeneración se transmite a través de esas dimensiones de degeneración para alcanzar la forma de salida. La semántica se describe en detalle en la página de emisión.

Call

Consulta también XlaBuilder::Call.

Invoca un procesamiento con los argumentos dados.

Call(computation, args...)

Argumentos Tipo Semántica
computation XlaComputation Cálculo de tipo T_0, T_1, ..., T_{N-1} -> S con N parámetros de tipo arbitrario
args secuencia de N XlaOp N argumentos de tipo arbitrario

La arity y los tipos de args deben coincidir con los parámetros de computation. No puede tener args.

Colesky

Consulta también XlaBuilder::Cholesky.

Calcula la descomposición de Colesky de un lote de matrices definidas positivas simétricas (hermitas).

Cholesky(a, lower)

Argumentos Tipo Semántica
a XlaOp un array de rango > 2 de un tipo complejo o de punto flotante.
lower bool si se debe usar el triángulo superior o inferior de a.

Si lower es true, calcula las matrices triangulares inferiores l, de modo que $a = l . l^T$. Si lower es false, calcula las matrices triangulares superiores u de modo que\(a = u^T . u\).

Los datos de entrada solo se leen desde el triángulo inferior/superior de a, según el valor de lower. Se ignoran los valores del otro triángulo. Los datos de salida se muestran en el mismo triángulo; los valores en el otro triángulo se definen en la implementación y pueden ser de cualquier tipo.

Si el rango de a es mayor que 2, a se trata como un lote de matrices, en el que todas, excepto las 2 dimensiones menores, son dimensiones de lote.

Si a no es un resultado positivo simétrico (hermitiano) definido por la implementación.

Pinzas

Consulta también XlaBuilder::Clamp.

Sujeta un operando para que se encuentre dentro del rango entre un valor mínimo y un máximo.

Clamp(min, operand, max)

Argumentos Tipo Semántica
min XlaOp array de tipo T
operand XlaOp array de tipo T
max XlaOp array de tipo T

Con un operando y valores mínimos y máximos, muestra el operando si está en el rango entre el mínimo y el máximo; de lo contrario, muestra el valor mínimo si el operando está por debajo de este rango o el valor máximo si está por encima de este rango. Es decir, clamp(a, x, b) = min(max(a, x), b).

Las tres matrices deben tener la misma forma. De manera alternativa, como una forma restringida de transmisión, min o max pueden ser un escalar de tipo T.

Ejemplo con min y max escalares:

let operand: s32[3] = {-1, 5, 9};
let min: s32 = 0;
let max: s32 = 6;
==>
Clamp(min, operand, max) = s32[3]{0, 5, 6};

Contraer

Consulta también XlaBuilder::Collapse y la operación tf.reshape.

Contrae las dimensiones de un array en una sola dimensión.

Collapse(operand, dimensions)

Argumentos Tipo Semántica
operand XlaOp array de tipo T
dimensions Vector de int64 en orden y consecutivo de las dimensiones de T.

La contracción reemplaza al subconjunto determinado de las dimensiones del operando por una sola dimensión. Los argumentos de entrada son un arreglo arbitrario de tipo T y un vector de constante de tiempo de compilación de índices de dimensión. Los índices de dimensión deben estar en orden (números de dimensión bajos a altos) y un subconjunto consecutivos de las dimensiones de T. Por lo tanto, {0, 1, 2}, {0, 1} o {1, 2} son todos conjuntos de dimensiones válidos, pero {1, 0} o {0, 2} no lo son. Se reemplazan por una única dimensión nueva, en la misma posición en la secuencia de dimensiones en la que se reemplazan, por el tamaño de dimensión nuevo igual al producto de los tamaños de dimensión originales. El número de dimensión más bajo en dimensions es la dimensión que varía más lenta (la más importante) en el nido de bucles que contrae estas dimensiones, y el número de dimensión más alto varía con mayor rapidez (el más pequeño). Consulta el operador tf.reshape si se necesita un orden de contracción más general.

Por ejemplo, supongamos que v es un array de 24 elementos:

let v = f32[4x2x3] { { {10, 11, 12},  {15, 16, 17} },
{ {20, 21, 22},  {25, 26, 27} },
{ {30, 31, 32},  {35, 36, 37} },
{ {40, 41, 42},  {45, 46, 47} } };

// Collapse to a single dimension, leaving one dimension.
let v012 = Collapse(v, {0,1,2});
then v012 == f32[24] {10, 11, 12, 15, 16, 17,
20, 21, 22, 25, 26, 27,
30, 31, 32, 35, 36, 37,
40, 41, 42, 45, 46, 47};

// Collapse the two lower dimensions, leaving two dimensions.
let v01 = Collapse(v, {0,1});
then v01 == f32[4x6] { {10, 11, 12, 15, 16, 17},
{20, 21, 22, 25, 26, 27},
{30, 31, 32, 35, 36, 37},
{40, 41, 42, 45, 46, 47} };

// Collapse the two higher dimensions, leaving two dimensions.
let v12 = Collapse(v, {1,2});
then v12 == f32[8x3] { {10, 11, 12},
{15, 16, 17},
{20, 21, 22},
{25, 26, 27},
{30, 31, 32},
{35, 36, 37},
{40, 41, 42},
{45, 46, 47} };

CollectivePermute

Consulta también XlaBuilder::CollectivePermute.

CollectivePermute es una operación colectiva que envía y recibe datos entre réplicas.

CollectivePermute(operand, source_target_pairs)

Argumentos Tipo Semántica
operand XlaOp array de entrada de dimensión n
source_target_pairs Vector de <int64, int64> Una lista de pares (source_replica_id, target_replica_id). Para cada par, el operando se envía de la réplica de origen a la réplica de destino.

Ten en cuenta que existen las siguientes restricciones en el source_target_pair:

  • Dos pares cualesquiera no deben tener el mismo ID de réplica de destino ni el mismo ID de réplica de origen.
  • Si un ID de réplica no es un destino en ningún par, la salida de esa réplica es un tensor que consta de 0 con la misma forma que la entrada.

Concatenate

Consulta también XlaBuilder::ConcatInDim.

Concatenate compone un array a partir de múltiples operandos de array. El array tiene la misma clasificación que cada uno de los operandos de array de entrada (que deben tener el mismo rango) y contiene los argumentos en el orden en que se especificaron.

Concatenate(operands..., dimension)

Argumentos Tipo Semántica
operands secuencia de N XlaOp N matrices de tipo T con dimensiones [L0, L1, ...]. Requiere N >= 1.
dimension int64 Un valor en el intervalo [0, N) que nombra la dimensión que se concatena entre los operands.

A excepción de dimension, todas las dimensiones deben ser iguales. Esto se debe a que XLA no admite arrays irregulares. Además, ten en cuenta que los valores de rango 0 no se pueden concatenar (ya que es imposible nombrar la dimensión a lo largo de la cual se produce la concatenación).

Ejemplo unidimensional:

Concat({ {2, 3}, {4, 5}, {6, 7} }, 0)
>>> {2, 3, 4, 5, 6, 7}

Ejemplo bidimensional:

let a = {
{1, 2},
{3, 4},
{5, 6},
};
let b = {
{7, 8},
};
Concat({a, b}, 0)
>>> {
{1, 2},
{3, 4},
{5, 6},
{7, 8},
}

Diagrama:

Condicionales

Consulta también XlaBuilder::Conditional.

Conditional(pred, true_operand, true_computation, false_operand, false_computation)

Argumentos Tipo Semántica
pred XlaOp Escalar de tipo PRED
true_operand XlaOp Argumento de tipo \(T_0\)
true_computation XlaComputation XlaComputation de tipo \(T_0 \to S\)
false_operand XlaOp Argumento de tipo \(T_1\)
false_computation XlaComputation XlaComputation de tipo \(T_1 \to S\)

Ejecuta true_computation si pred es true, false_computation si pred es false, y muestra el resultado.

true_computation debe tener un solo argumento de tipo \(T_0\) y se invocará con true_operand, que debe ser del mismo tipo. false_computation debe tener un solo argumento de tipo \(T_1\) y se invocará con false_operand, que debe ser del mismo tipo. El tipo del valor que se muestra de true_computation y false_computation debe ser el mismo.

Ten en cuenta que solo se ejecutará uno de los valores true_computation y false_computation según el valor de pred.

Conditional(branch_index, branch_computations, branch_operands)

Argumentos Tipo Semántica
branch_index XlaOp Escalar de tipo S32
branch_computations secuencia de N XlaComputation XlaComputations de tipo \(T_0 \to S , T_1 \to S , ..., T_{N-1} \to S\)
branch_operands secuencia de N XlaOp Argumentos de tipo \(T_0 , T_1 , ..., T_{N-1}\)

Ejecuta branch_computations[branch_index] y muestra el resultado. Si branch_index es una S32 que es < 0 o >= N, entonces branch_computations[N-1] se ejecuta como la rama predeterminada.

Cada branch_computations[b] debe tener un solo argumento de tipo \(T_b\) y se invocará con branch_operands[b], que debe ser del mismo tipo. El tipo del valor mostrado de cada branch_computations[b] debe ser el mismo.

Ten en cuenta que solo se ejecutará una de las branch_computations según el valor de branch_index.

Conv. (convolución)

Consulta también XlaBuilder::Conv.

Como ConvWithGeneralPadding, pero el padding se especifica de forma abreviada como SAME o VALID. SAME rellena la entrada (lhs) con ceros para que la salida tenga la misma forma que la entrada cuando no se tienen en cuenta los pasos. Si hay relleno VALID, simplemente significa que no hay relleno.

ConvConGeneralPadding (convolución)

Consulta también XlaBuilder::ConvWithGeneralPadding.

Calcula una convolución del tipo usado en las redes neuronales. Aquí, una convolución puede considerarse como una ventana de n dimensiones que se mueve por un área de base de esa dimensión, y se realiza un cálculo para cada posición posible de la ventana.

Argumentos Tipo Semántica
lhs XlaOp Array de entradas rango n+2
rhs XlaOp matricial n+2 de pesos del kernel
window_strides ArraySlice<int64> Array n-d de segmentos de kernel
padding ArraySlice< pair<int64,int64>> Array n-d de relleno (bajo, alto)
lhs_dilation ArraySlice<int64> Array de factor de dilatación de n-d lhs
rhs_dilation ArraySlice<int64> Array de factor de dilatación de n-d rhs
feature_group_count int64 la cantidad de grupos de atributos
batch_group_count int64 la cantidad de grupos por lotes

Supongamos que n es el número de dimensiones espaciales. El argumento lhs es un array de clasificación n+2 que describe el área de la base. Esto se llama entrada, aunque, por supuesto, la RHS también es una entrada. En una red neuronal, estas son las activaciones de entrada. Las dimensiones n+2 están en este orden:

  • batch: Cada coordenada en esta dimensión representa una entrada independiente para la que se realiza la convolución.
  • z/depth/features: Cada posición (y,x) en el área de la base tiene un vector asociado, que entra en esta dimensión.
  • spatial_dims: Describe las dimensiones espaciales n que definen el área base por la que se mueve la ventana.

El argumento rhs es un array de rango n+2 que describe el filtro, kernel o ventana convolucional. Las dimensiones se muestran en este orden:

  • output-z: Es la dimensión z del resultado.
  • input-z: El tamaño de esta dimensión multiplicado por feature_group_count debe igual al tamaño de la dimensión z en lhs.
  • spatial_dims: Describe las dimensiones espaciales n que definen la ventana n-d que se mueve por el área base.

El argumento window_strides especifica el segmento de la ventana convolucional en las dimensiones espaciales. Por ejemplo, si el segmento de la primera dimensión espacial es 3, la ventana solo se puede colocar en las coordenadas en las que el primer índice espacial es divisible por 3.

El argumento padding especifica la cantidad de padding cero que se aplicará al área de la base. La cantidad de padding puede ser negativa; el valor absoluto del padding negativo indica la cantidad de elementos que se quitarán de la dimensión especificada antes de realizar la convolución. padding[0] especifica el relleno de la dimensión y y padding[1] especifica el relleno de la dimensión x. Cada par tiene el padding bajo como primer elemento y el padding alto como el segundo. El padding bajo se aplica en la dirección de los índices más bajos, mientras que el padding alto se aplica en la dirección de los índices más altos. Por ejemplo, si padding[1] es (2,3), habrá un padding de 2 ceros a la izquierda y 3 ceros a la derecha en la segunda dimensión espacial. El uso de relleno es equivalente a insertar esos mismos valores de cero en la entrada (lhs) antes de realizar la convolución.

Los argumentos lhs_dilation y rhs_dilation especifican el factor de dilatación que se aplicará a las lhs y rhs, respectivamente, en cada dimensión espacial. Si el factor de dilatación en una dimensión espacial es d, se colocan los agujeros de d-1 de forma implícita entre cada una de las entradas de esa dimensión, lo que aumenta el tamaño del array. Los agujeros se llenan con un valor no-ops, que para la convolución significa ceros.

La dilatación de la RHS también se conoce como convolución atros. Si deseas obtener información más detallada, consulta tf.nn.atrous_conv2d. La dilatación de las lhs también se denomina convolución transpuesta. Para obtener información detallada, consulta tf.nn.conv2d_transpose.

El argumento feature_group_count (valor predeterminado 1) se puede usar para las convoluciones agrupadas. feature_group_count debe ser un divisor de la dimensión del atributo de entrada y de la de salida. Si feature_group_count es mayor que 1, significa que, conceptualmente, la dimensión del atributo de entrada y salida y la dimensión del atributo de salida rhs se dividen de manera uniforme en muchos grupos de feature_group_count, cada uno de los cuales consta de una subsecuencia consecutiva de atributos. La dimensión del atributo de entrada de rhs debe ser igual a la dimensión del atributo de entrada lhs dividida por feature_group_count (por lo que ya tiene el tamaño de un grupo de atributos de entrada). Los grupos i-th se usan juntos para calcular feature_group_count para muchas convoluciones distintas. Los resultados de estas convoluciones se concatenan en la dimensión del atributo de salida.

Para la convolución en profundidad, el argumento feature_group_count se establecería en la dimensión del atributo de entrada y el filtro se cambiaría de [filter_height, filter_width, in_channels, channel_multiplier] a [filter_height, filter_width, 1, in_channels * channel_multiplier]. Para obtener más información, consulta tf.nn.depthwise_conv2d.

El argumento batch_group_count (valor predeterminado 1) se puede usar para los filtros agrupados durante la propagación inversa. batch_group_count debe ser un divisor del tamaño de la dimensión del lote lhs (entrada). Si batch_group_count es mayor que 1, significa que la dimensión del lote de salida debe ser de tamaño input batch / batch_group_count. El valor de batch_group_count debe ser un divisor del tamaño del atributo de salida.

La forma de salida tiene las siguientes dimensiones, en este orden:

  • batch: El tamaño de esta dimensión multiplicado por batch_group_count debe ser igual al tamaño de la dimensión batch en lhs.
  • z: Tiene el mismo tamaño que output-z en el kernel (rhs).
  • spatial_dims: Un valor para cada posición válida de la ventana convolucional

En la figura anterior, se muestra cómo funciona el campo batch_group_count. En efecto, dividimos cada lote en grupos batch_group_count y hacemos lo mismo para los atributos de salida. Luego, para cada uno de estos grupos, realizamos convoluciones en pares y concatenamos la salida junto con la dimensión del atributo de salida. La semántica operativa de todas las demás dimensiones (atributo y espacial) sigue siendo la misma.

Las posiciones válidas de la ventana convolucional se determinan por los pasos y el tamaño del área de la base después del relleno.

Para describir lo que hace una convolución, considera una convolución en 2D y elige algunas coordenadas fijas batch, z, y y x en el resultado. Entonces, (y,x) es una posición de una esquina de la ventana dentro del área de la base (p.ej., la esquina superior izquierda, según cómo interpretes las dimensiones espaciales). Ahora tenemos una ventana de 2D, tomada de la zona de la base, en la que cada punto 2D está asociado a un vector 1d, por lo que obtenemos un cuadro de 3D. Desde el kernel convolucional, desde que corregimos la coordenada de salida z, también tenemos un cuadro 3D. Los dos cuadros tienen las mismas dimensiones, por lo que podemos tomar la suma de los productos a nivel de elementos entre los dos cuadros (similar a un producto escalar). Ese es el valor de salida.

Ten en cuenta que si output-z es p.ej., 5, cada posición de la ventana produce 5 valores en el resultado en la dimensión z del resultado. Estos valores difieren en qué parte del kernel convolucional se usa: hay un cuadro de valores 3D separado para cada coordenada output-z. Puedes pensar en ello como 5 convoluciones independientes con un filtro diferente para cada una.

A continuación, se muestra el seudocódigo de una convolución en 2D con relleno y pasos:

for (b, oz, oy, ox) {  // output coordinates
  value = 0;
  for (iz, ky, kx) {  // kernel coordinates and input z
    iy = oy*stride_y + ky - pad_low_y;
    ix = ox*stride_x + kx - pad_low_x;
    if ((iy, ix) inside the base area considered without padding) {
      value += input(b, iz, iy, ix) * kernel(oz, iz, ky, kx);
    }
  }
  output(b, oz, oy, ox) = value;
}

ConvertElementType

Consulta también XlaBuilder::ConvertElementType.

Al igual que static_cast en C++, realiza una operación de conversión a nivel de elementos de una forma de datos a una forma objetivo. Las dimensiones deben coincidir, y la conversión depende de los elementos; p.ej., los elementos s32 se convierten en elementos f32 mediante una rutina de conversión de s32 a f32.

ConvertElementType(operand, new_element_type)

Argumentos Tipo Semántica
operand XlaOp matriz de tipo T con atenuación D
new_element_type PrimitiveType tipo U

Las dimensiones del operando y de la forma objetivo deben coincidir. Los tipos de elementos de origen y de destino no deben ser tuplas.

Una conversión como T=s32 a U=f32 realizará una rutina de conversión de valor de int a float que se normaliza, como el redondeo a par más cercano.

let a: s32[3] = {0, 1, 2};
let b: f32[3] = convert(a, f32);
then b == f32[3]{0.0, 1.0, 2.0}

CrossReplicaSum

Realiza AllReduce con un cálculo de suma.

CustomCall

Consulta también XlaBuilder::CustomCall.

Llama a una función proporcionada por el usuario dentro de un cálculo.

CustomCall(target_name, args..., shape)

Argumentos Tipo Semántica
target_name string Nombre de la función. Se emitirá una instrucción de llamada orientada a este nombre de símbolo.
args secuencia de N XlaOp N argumentos de tipo arbitrario, que se pasarán a la función.
shape Shape Forma del resultado de la función

La firma de la función es la misma, independientemente de la aridad o el tipo de argumentos:

extern "C" void target_name(void* out, void** in);

Por ejemplo, si CustomCall se usa de la siguiente manera:

let x = f32[2] {1,2};
let y = f32[2x3] { {10, 20, 30}, {40, 50, 60} };

CustomCall("myfunc", {x, y}, f32[3x3])

Este es un ejemplo de una implementación de myfunc:

extern "C" void myfunc(void* out, void** in) {
  float (&x)[2] = *static_cast<float(*)[2]>(in[0]);
  float (&y)[2][3] = *static_cast<float(*)[2][3]>(in[1]);
  EXPECT_EQ(1, x[0]);
  EXPECT_EQ(2, x[1]);
  EXPECT_EQ(10, y[0][0]);
  EXPECT_EQ(20, y[0][1]);
  EXPECT_EQ(30, y[0][2]);
  EXPECT_EQ(40, y[1][0]);
  EXPECT_EQ(50, y[1][1]);
  EXPECT_EQ(60, y[1][2]);
  float (&z)[3][3] = *static_cast<float(*)[3][3]>(out);
  z[0][0] = x[1] + y[1][0];
  // ...
}

La función proporcionada por el usuario no debe tener efectos secundarios y su ejecución debe ser idempotente.

Punto

Consulta también XlaBuilder::Dot.

Dot(lhs, rhs)

Argumentos Tipo Semántica
lhs XlaOp array de tipo T
rhs XlaOp array de tipo T

La semántica exacta de esta operación depende de las clasificaciones de los operandos:

Entrada Resultado Semántica
vector [n] dot vector [n] escalar producto vectorial de puntos
matriz [m x k] dot vector [k] vector [m] multiplicación de vectores de matrices
matriz [m x k] dot matriz [k x n] matriz [m x n] multiplicación de matrices y matrices

La operación realiza la suma de productos en la segunda dimensión de lhs (o el primero si tiene clasificación 1) y la primera dimensión de rhs. Estas son las dimensiones "contratadas". Las dimensiones contraídas de lhs y rhs deben ser del mismo tamaño. En la práctica, se puede usar para realizar productos de puntos entre vectores, multiplicaciones de vectores o matrices o multiplicaciones de matrices y matrices.

DotGeneral

Consulta también XlaBuilder::DotGeneral.

DotGeneral(lhs, rhs, dimension_numbers)

Argumentos Tipo Semántica
lhs XlaOp array de tipo T
rhs XlaOp array de tipo T
dimension_numbers DotDimensionNumbers números de dimensión de lote y contratación

Es similar a Dot, pero permite que se especifiquen números de dimensión de lote y de contratación para lhs y rhs.

Campos DotDimensionNumbers Tipo Semántica
lhs_contracting_dimensions int64 repetido Números de dimensión de contratación lhs
rhs_contracting_dimensions int64 repetido Números de dimensión de contratación rhs
lhs_batch_dimensions int64 repetido lhs de números de dimensión del lote
rhs_batch_dimensions int64 repetido rhs de números de dimensión del lote

DotGeneral realiza la suma de productos sobre las dimensiones de contratación especificadas en dimension_numbers.

No es necesario que los números de dimensión de contratación asociados de lhs y rhs sean iguales, pero deben tener los mismos tamaños de dimensión.

Ejemplo con números de dimensión de contratación:

lhs = { {1.0, 2.0, 3.0},
{4.0, 5.0, 6.0} }

rhs = { {1.0, 1.0, 1.0},
{2.0, 2.0, 2.0} }

DotDimensionNumbers dnums;
dnums.add_lhs_contracting_dimensions(1);
dnums.add_rhs_contracting_dimensions(1);

DotGeneral(lhs, rhs, dnums) -> { {6.0, 12.0},
{15.0, 30.0} }

Los números de dimensión del lote asociados de lhs y rhs deben tener los mismos tamaños de dimensión.

Ejemplo con números de dimensión del lote (tamaño del lote 2, matrices 2x2):

lhs = { { {1.0, 2.0},
{3.0, 4.0} },
{ {5.0, 6.0},
{7.0, 8.0} } }

rhs = { { {1.0, 0.0},
{0.0, 1.0} },
{ {1.0, 0.0},
{0.0, 1.0} } }

DotDimensionNumbers dnums;
dnums.add_lhs_contracting_dimensions(2);
dnums.add_rhs_contracting_dimensions(1);
dnums.add_lhs_batch_dimensions(0);
dnums.add_rhs_batch_dimensions(0);

DotGeneral(lhs, rhs, dnums) -> { { {1.0, 2.0},
{3.0, 4.0} },
{ {5.0, 6.0},
{7.0, 8.0} } }
Entrada Resultado Semántica
[b0, m, k] dot [b0, k, n] [b0, m, n] matmul por lotes
[b0, b1, m, k] dot [b0, b1, k y n] [b0, b1, m y n]. matmul por lotes

A continuación, el número de dimensión resultante comienza con la dimensión del lote, luego, la dimensión lhs no contraída/no por lotes y, por último, rhs la dimensión no contraído/no por lotes.

DynamicSlice

Consulta también XlaBuilder::DynamicSlice.

DynamicSlice extrae un subarreglo del arreglo de entrada en start_indices dinámico. El tamaño de la porción en cada dimensión se pasa en size_indices, que especifica el punto de finalización de los intervalos de porción exclusivos en cada dimensión: [inicio, inicio + tamaño). La forma de start_indices debe tener una clasificación == 1, con un tamaño de dimensión igual al rango de operand.

DynamicSlice(operand, start_indices, size_indices)

Argumentos Tipo Semántica
operand XlaOp Array N dimensional de tipo T
start_indices secuencia de N XlaOp Lista de N números enteros escalares que contienen los índices iniciales de la porción para cada dimensión. El valor debe ser mayor o igual que cero.
size_indices ArraySlice<int64> Lista de N números enteros que contienen el tamaño de porción de cada dimensión. Cada valor debe ser estrictamente mayor que cero, y el valor de inicio + tamaño debe ser menor o igual que el tamaño de la dimensión para evitar ajustar el tamaño de la dimensión del módulo.

Los índices de porción eficaces se calculan mediante la aplicación de la siguiente transformación para cada índice i en [1, N) antes de realizar la porción:

start_indices[i] = clamp(start_indices[i], 0, operand.dimension_size[i] - size_indices[i])

Esto garantiza que la porción extraída siempre esté dentro de los límites con respecto al arreglo del operando. Si la porción está dentro de los límites antes de que se aplique la transformación, la transformación no tiene efecto.

Ejemplo unidimensional:

let a = {0.0, 1.0, 2.0, 3.0, 4.0}
let s = {2}

DynamicSlice(a, s, {2}) produces:
{2.0, 3.0}

Ejemplo bidimensional:

let b =
{ {0.0,  1.0,  2.0},
{3.0,  4.0,  5.0},
{6.0,  7.0,  8.0},
{9.0, 10.0, 11.0} }
let s = {2, 1}

DynamicSlice(b, s, {2, 2}) produces:
{ { 7.0,  8.0},
{10.0, 11.0} }

DynamicUpdateSlice

Consulta también XlaBuilder::DynamicUpdateSlice.

DynamicUpdateSlice genera un resultado que es el valor del arreglo de entrada operand, con una porción update reemplazada en start_indices. La forma de update determina la forma del subarreglo del resultado que se actualiza. La forma de start_indices debe tener una clasificación == 1, con un tamaño de dimensión igual al rango de operand.

DynamicUpdateSlice(operand, update, start_indices)

Argumentos Tipo Semántica
operand XlaOp Array N dimensional de tipo T
update XlaOp Array de dimensión N de tipo T que contiene la actualización de porción. Cada dimensión de la forma de actualización debe ser estrictamente mayor que cero, y el valor de inicio + actualización debe ser menor o igual que el tamaño del operando de cada dimensión para evitar generar índices de actualización fuera de los límites.
start_indices secuencia de N XlaOp Lista de N números enteros escalares que contienen los índices iniciales de la porción para cada dimensión. El valor debe ser mayor o igual que cero.

Los índices de porción eficaces se calculan mediante la aplicación de la siguiente transformación para cada índice i en [1, N) antes de realizar la porción:

start_indices[i] = clamp(start_indices[i], 0, operand.dimension_size[i] - update.dimension_size[i])

Esto garantiza que la porción actualizada siempre esté dentro de los límites del arreglo de operando. Si la porción está dentro de los límites antes de que se aplique la transformación, la transformación no tiene efecto.

Ejemplo unidimensional:

let a = {0.0, 1.0, 2.0, 3.0, 4.0}
let u = {5.0, 6.0}
let s = {2}

DynamicUpdateSlice(a, u, s) produces:
{0.0, 1.0, 5.0, 6.0, 4.0}

Ejemplo bidimensional:

let b =
{ {0.0,  1.0,  2.0},
{3.0,  4.0,  5.0},
{6.0,  7.0,  8.0},
{9.0, 10.0, 11.0} }
let u =
{ {12.0,  13.0},
{14.0,  15.0},
{16.0,  17.0} }

let s = {1, 1}

DynamicUpdateSlice(b, u, s) produces:
{ {0.0,  1.0,  2.0},
{3.0, 12.0, 13.0},
{6.0, 14.0, 15.0},
{9.0, 16.0, 17.0} }

Operaciones aritméticas binarias en elementos

Consulta también XlaBuilder::Add.

Se admite un conjunto de operaciones aritméticas binarias de elementos.

Op(lhs, rhs)

Donde Op es Add (suma), Sub (resta), Mul (multiplicación), Div (división), Rem (resto), Max (máximo), Min (mínimo), LogicalAnd (operación lógica AND) o LogicalOr (operación lógica OR).

Argumentos Tipo Semántica
lhs XlaOp operando izquierdo: array de tipo T
rhs XlaOp operando derecho: array de tipo T

Las formas de los argumentos deben ser similares o compatibles. Consulta la documentación de transmisión sobre lo que significa que las formas sean compatibles. El resultado de una operación tiene una forma que es el resultado de transmitir los dos arrays de entrada. En esta variante, las operaciones entre arreglos de diferentes clasificaciones no son compatibles, a menos que uno de los operandos sea un escalar.

Cuando Op es Rem, el signo del resultado se toma del dividendo, y el valor absoluto del resultado es siempre menor que el valor absoluto del divisor.

El desbordamiento de la división de enteros (división o resto con signo o sin firma por cero o división o resto con signo de INT_SMIN con -1) produce un valor definido de implementación.

Existe una variante alternativa con compatibilidad de transmisión de diferente rango para estas operaciones:

Op(lhs, rhs, broadcast_dimensions)

En el ejemplo anterior, Op es igual al anterior. Esta variante de la operación se debe usar para operaciones aritméticas entre arrays de diferentes rangos (por ejemplo, agregar una matriz a un vector).

El operando broadcast_dimensions adicional es una porción de números enteros que se usa para expandir el rango del operando de rango inferior hasta el rango del operando de rango superior. broadcast_dimensions asigna las dimensiones de la forma de rango inferior a las dimensiones de la forma de rango superior. Las dimensiones sin asignar de la forma expandida se rellenan con dimensiones de tamaño uno. Luego, la transmisión de la dimensión de degeneración transmite las formas a lo largo de estas dimensiones de degeneración para igualar las formas de ambos operandos. La semántica se describe en detalle en la página de emisión.

Operaciones de comparación en elementos

Consulta también XlaBuilder::Eq.

Se admite un conjunto de operaciones estándar de comparación binaria a nivel de elementos. Ten en cuenta que se aplica la semántica de comparación de punto flotante estándar IEEE 754 cuando se comparan tipos de punto flotante.

Op(lhs, rhs)

Donde Op es Eq (igual a), Ne (no igual a), Ge (mayor o igual que), Gt (mayor que), Le (menor o igual que), Lt (menor que). Otro conjunto de operadores, EqTotalOrder, NeTotalOrder, GeTotalOrder, GtTotalOrder, LeTotalOrder y LtTotalOrder, además admiten un pedido total sobre los números de punto flotante, mediante la aplicación -NaN < -Inf < -Finite < -0 < +N < +Finite <a +N <N

Argumentos Tipo Semántica
lhs XlaOp operando izquierdo: array de tipo T
rhs XlaOp operando derecho: array de tipo T

Las formas de los argumentos deben ser similares o compatibles. Consulta la documentación de transmisión sobre lo que significa que las formas sean compatibles. El resultado de una operación tiene una forma que es el resultado de transmitir los dos arrays de entrada con el tipo de elemento PRED. En esta variante, no se admiten las operaciones entre arreglos de diferentes clasificaciones, a menos que uno de los operandos sea un escalar.

Existe una variante alternativa con compatibilidad de transmisión de diferente rango para estas operaciones:

Op(lhs, rhs, broadcast_dimensions)

En el ejemplo anterior, Op es igual al anterior. Esta variante de la operación se debe usar para operaciones de comparación entre arrays de diferentes clasificaciones (por ejemplo, agregar una matriz a un vector).

El operando broadcast_dimensions adicional es una porción de números enteros que especifican las dimensiones que se usarán para transmitir los operandos. La semántica se describe en detalle en la página de emisión.

Funciones unarias de elementos

XlaBuilder admite estas funciones unarias a nivel de elementos:

Abs(operand) Abs. de elementos x -> |x|.

Ceil(operand) Cielo por elementos x -> ⌈x⌉.

Cos(operand) Coseno por elemento x -> cos(x).

Exp(operand) x -> e^x exponencial natural por elementos.

Floor(operand) Plano básico x -> ⌊x⌋.

Imag(operand) Parte imaginaria de elementos de una forma compleja (o real). x -> imag(x). Si el operando es un tipo de punto flotante, muestra 0.

IsFinite(operand) Comprueba si cada elemento de operand es finito, es decir, no es infinito positivo ni negativo, y no es NaN. Muestra un array de valores PRED con la misma forma que la entrada, en el que cada elemento es true solo si el elemento de entrada correspondiente es finito.

Log(operand) Logaritmo natural por elementos x -> ln(x).

LogicalNot(operand) Lógico de elemento, no x -> !(x).

Logistic(operand) Cálculo de la función logística por elementos x -> logistic(x).

PopulationCount(operand) Calcula la cantidad de bits establecidos en cada elemento de operand.

Neg(operand) Negación por elementos x -> -x.

Real(operand) Parte real por elementos de una forma compleja (o real). x -> real(x). Si el operando es un tipo de punto flotante, muestra el mismo valor.

Rsqrt(operand) Recíproco a nivel de elemento de la operación de raíz cuadrada x -> 1.0 / sqrt(x).

Sign(operand) Operación de signo por elementos x -> sgn(x), donde

\[\text{sgn}(x) = \begin{cases} -1 & x < 0\\ -0 & x = -0\\ NaN & x = NaN\\ +0 & x = +0\\ 1 & x > 0 \end{cases}\]

Usando el operador de comparación del tipo de elemento de operand.

Sqrt(operand) Operación de raíz cuadrada por elementos x -> sqrt(x).

Cbrt(operand) Operación de raíz cúbica por elementos x -> cbrt(x).

Tanh(operand) Tangente hiperbólica con elementos x -> tanh(x).

Round(operand) Redondeo por elementos, empate desde cero.

RoundNearestEven(operand) Redondeo por elementos, se vincula al par más cercano.

Argumentos Tipo Semántica
operand XlaOp El operando de la función

La función se aplica a cada elemento del array operand, lo que da como resultado un array con la misma forma. Se permite que operand sea un escalar (rango 0).

Fft

La operación FFT de XLA implementa las Transformaciones de Fourier inversas y futuras para entradas y salidas reales y complejas. Se admiten FFT multidimensionales en hasta 3 ejes.

Consulta también XlaBuilder::Fft.

Argumentos Tipo Semántica
operand XlaOp La matriz que estamos transformando Fourier.
fft_type FftType Consulta la tabla que se encuentra a continuación
fft_length ArraySlice<int64> Las longitudes de dominio de tiempo de los ejes que se transforman. Esto es necesario en particular para que IRFFT redimensione el eje más interno, ya que RFFT(fft_length=[16]) tiene la misma forma de salida que RFFT(fft_length=[17]).
FftType Semántica
FFT Reenvía los FFT complejos a complejos. La forma no se modifica.
IFFT FFT inverso de complejo a complejo. La forma no se modifica.
RFFT Reenvía los FFT reales a complejos. La forma del eje más interno se reduce a fft_length[-1] // 2 + 1 si fft_length[-1] es un valor distinto de cero, y se omite la parte conjugada invertida de la señal transformada más allá de la frecuencia de Nyquist.
IRFFT FFT inverso de real a complejo (es decir, toma complejo, muestra real). La forma del eje más interno se expande a fft_length[-1] si fft_length[-1] es un valor distinto de cero, lo que infiere la parte de la señal transformada más allá de la frecuencia de Nyquist a partir del conjugado inverso de las entradas 1 a fft_length[-1] // 2 + 1.

FFT multidimensional

Cuando se proporciona más de 1 fft_length, esto equivale a aplicar una cascada de operaciones de FFT a cada uno de los ejes más internos. Ten en cuenta que, en los casos reales->complejos y complejos->reales, la transformación del eje interno se realiza primero (RFFT; el último para IRFFT), por lo que el eje más interno es el que cambia de tamaño. Por lo tanto, otras transformaciones de ejes serán complejas.

Detalles de la implementación

Los FFT de CPU cuentan con el respaldo de TensorFFT de Eigen. GPU FFT usa cuFFT.

Gather

La operación de recopilación de XLA une varias porciones (cada una en un desplazamiento de tiempo de ejecución potencialmente diferente) de un array de entrada.

Semántica general

Consulta también XlaBuilder::Gather. Para obtener una descripción más intuitiva, consulta la sección "Descripción informal" que aparece a continuación.

gather(operand, start_indices, offset_dims, collapsed_slice_dims, slice_sizes, start_index_map)

Argumentos Tipo Semántica
operand XlaOp El array del cual recopilamos.
start_indices XlaOp Arreglo que contiene los índices iniciales de las porciones que recopilamos.
index_vector_dim int64 Es la dimensión de start_indices que "contiene" los índices iniciales. Obtén una descripción detallada a continuación.
offset_dims ArraySlice<int64> Conjunto de dimensiones en la forma de salida que se desplaza en un arreglo dividido del operando.
slice_sizes ArraySlice<int64> slice_sizes[i] son los límites de la porción en la dimensión i.
collapsed_slice_dims ArraySlice<int64> Es el conjunto de dimensiones de cada porción que se contraen. Estas dimensiones deben ser de tamaño 1.
start_index_map ArraySlice<int64> Un mapa en el que se describe cómo asignar índices en start_indices a índices legales en un operando.
indices_are_sorted bool Indica si se garantiza que los índices estén ordenados por el emisor.

Para mayor comodidad, etiquetamos las dimensiones en el array de salida que no están en offset_dims como batch_dims.

El resultado es un array de clasificación batch_dims.size + offset_dims.size.

El valor de operand.rank debe ser igual a la suma de offset_dims.size y collapsed_slice_dims.size. Además, slice_sizes.size debe ser igual a operand.rank.

Si index_vector_dim es igual a start_indices.rank, consideramos implícitamente que start_indices tenga una dimensión 1 final (es decir, si start_indices era de forma [6,7] y index_vector_dim es 2, consideramos implícitamente que la forma de start_indices sea [6,7,1]).

Los límites del array de salida en la dimensión i se calculan de la siguiente manera:

  1. Si i está presente en batch_dims (es decir, es igual a batch_dims[k] para algunos k), elegimos los límites de dimensión correspondientes de start_indices.shape y se omite index_vector_dim (es decir, selecciona start_indices.shape.dims[k] si k < index_vector_dim y start_indices.shape.dims[k+1]).

  2. Si i está presente en offset_dims (es decir, igual a offset_dims[k] para algunos k), elegiremos el límite correspondiente de slice_sizes después de considerar collapsed_slice_dims (es decir, seleccionamos adjusted_slice_sizes[k], donde adjusted_slice_sizes es slice_sizes con los límites de los índices collapsed_slice_dims quitados).

De manera formal, el índice de operando In correspondiente a un índice de salida Out determinado se calcula de la siguiente manera:

  1. Deja que G = { Out[k] para k en batch_dims }. Usa G para dividir un vector S de manera que S[i] = start_indices[Combine(G, i)], en el que Combine(A, b) inserta b en la posición index_vector_dim en A. Ten en cuenta que esto está bien definido, incluso si G está vacío: si G está vacío, entonces S = start_indices.

  2. Crea un índice inicial, Sin, en operand mediante S mediante la dispersión de S con start_index_map. Más precisamente:

    1. Sin[start_index_map[k]] = S[k] si k < start_index_map.size.

    2. Sin[_] = 0 de lo contrario.

  3. Crea un índice Oin en operand mediante la dispersión de los índices en las dimensiones de desplazamiento en Out de acuerdo con el conjunto collapsed_slice_dims. Más precisamente:

    1. Oin[remapped_offset_dims(k)] = Out[offset_dims[k]] si k < offset_dims.size (remapped_offset_dims se define a continuación).

    2. Oin[_] = 0 de lo contrario.

  4. In es Oin + Sin, donde + es la suma de los elementos.

remapped_offset_dims es una función monótona con dominio [0, offset_dims.size) y rango [0, operand.rank) \ collapsed_slice_dims. Si, p.ej., offset_dims.size es 4, operand.rank es 6 y collapsed_slice_dims es {0, 2}, luego, remapped_offset_dims es {01, 13, 24, 35}.

Si indices_are_sorted se configura como verdadero, XLA puede suponer que el usuario ordena start_indices (en orden ascendente start_index_map). Si no lo están, entonces se define la semántica de la implementación.

Descripción informal y ejemplos

Informalmente, cada índice Out en el array de salida corresponde a un elemento E en el array de operando, y se calcula de la siguiente manera:

  • Usamos las dimensiones del lote en Out para buscar un índice de inicio a partir de start_indices.

  • Usamos start_index_map para asignar el índice de inicio (cuyo tamaño puede ser menor que operand.rank) a un índice de inicio "completo" en operand.

  • Dividimos de forma dinámica una porción con el tamaño slice_sizes usando el índice de inicio completo.

  • Para cambiar la forma de la porción, contrae las dimensiones de collapsed_slice_dims. Dado que todas las dimensiones de porciones contraídas deben tener un límite de 1, este cambio de forma siempre es legal.

  • Usamos las dimensiones de desplazamiento en Out para indexar en esta porción y obtener el elemento de entrada, E, que corresponde al índice de salida Out.

index_vector_dim se configura como start_indices.rank - 1 en todos los ejemplos siguientes. Los valores más interesantes para index_vector_dim no cambian la operación en esencia, pero hacen que la representación visual sea más engorrosa.

Para tener una intuición de cómo se combina todo lo anterior, echemos un vistazo a un ejemplo que reúne 5 porciones de forma [8,6] de un arreglo [16,11]. La posición de una porción en el array [16,11] se puede representar como un vector de índice de la forma S64[2], por lo que el conjunto de 5 posiciones se puede representar como un array S64[5,2].

El comportamiento de la operación de recopilación se puede representar como una transformación de índices que toma [G,O0,O1], un índice en la forma de salida y lo asigna a un elemento en el array de entrada de la siguiente manera:

Primero, seleccionamos un vector (X, Y) del array de recopilación de índices con G. El elemento del array de salida en el índice [G,O0,O1] es el elemento en el array de entrada en el índice [X + O0,Y+O1].

slice_sizes es [8,6], que decide el rango de O0 y O1, y esto, a su vez, decide los límites de la porción.

Esta operación de recopilación actúa como una porción dinámica por lotes con G como la dimensión del lote.

Los índices de recopilación pueden ser multidimensionales. Por ejemplo, una versión más general del ejemplo anterior que use un array de "índices de recopilación" con la forma [4,5,2] traduciría índices como el siguiente:

Nuevamente, esto actúa como una porción dinámica por lotes G0 y G1 como las dimensiones del lote. El tamaño de la porción sigue siendo [8,6].

La operación de recopilación en XLA generaliza la semántica informal descrita anteriormente de las siguientes maneras:

  1. Podemos configurar qué dimensiones de la forma de salida son las de desplazamiento (las que contienen O0, O1 en el último ejemplo). Las dimensiones del lote de salida (dimensiones que contienen G0, G1 en el último ejemplo) se definen como las dimensiones de salida que no son dimensiones de desplazamiento.

  2. La cantidad de dimensiones de desplazamiento de salida presentes explícitamente en la forma de salida puede ser menor que la clasificación de entrada. Estas dimensiones "faltantes", que se enumeran explícitamente como collapsed_slice_dims, deben tener un tamaño de porción de 1. Dado que tienen un tamaño de porción de 1, el único índice válido para ellos es 0, y si los omites, no se genera ambigüedad.

  3. La porción extraída del arreglo “Reúne índices” (X, Y) en el último ejemplo puede tener menos elementos que la clasificación del arreglo de entrada, y una asignación explícita determina cómo se debe expandir el índice para que tenga la misma clasificación que la entrada.

Como ejemplo final, usamos (2) y (3) para implementar tf.gather_nd:

G0 y G1 se usan para dividir un índice de inicio del arreglo de índices de recopilación como de costumbre, excepto que el índice inicial tiene un solo elemento, X. Del mismo modo, solo hay un índice de desplazamiento de salida con el valor O0. Sin embargo, antes de usarse como índices en el array de entrada, estos se expanden de acuerdo con "Recopilar asignación de índices" (start_index_map en la descripción formal) y "Asignación de desplazamiento" (remapped_offset_dims en la descripción formal) a [X,0] y [0,O0], respectivamente, y se agregan hasta [X,O0]. En otras palabras, el índice de entrada [G0].G para O la salida [G]. G0000O11GatherIndicestf.gather_nd

slice_sizes para este caso es [1,11]. De forma intuitiva, esto significa que cada índice X en el arreglo de índices de conjunto elige una fila completa y el resultado es la concatenación de todas estas filas.

GetDimensionSize

Consulta también XlaBuilder::GetDimensionSize.

Muestra el tamaño de la dimensión determinada del operando. El operando debe tener forma de array.

GetDimensionSize(operand, dimension)

Argumentos Tipo Semántica
operand XlaOp array de entrada de dimensión n
dimension int64 Un valor en el intervalo [0, n) que especifica la dimensión

SetDimensionSize

Consulta también XlaBuilder::SetDimensionSize.

Establece el tamaño dinámico de la dimensión determinada de XlaOp. El operando debe tener forma de array.

SetDimensionSize(operand, size, dimension)

Argumentos Tipo Semántica
operand XlaOp de entrada n.
size XlaOp int32 que representa el tamaño dinámico del entorno de ejecución
dimension int64 Es un valor en el intervalo [0, n) que especifica la dimensión.

Pasa el operando como resultado, con la dimensión dinámica que realiza el seguimiento del compilador.

Las operaciones de reducción descendentes ignorarán los valores con relleno.

let v: f32[10] = f32[10]{1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
let five: s32 = 5;
let six: s32 = 6;

// Setting dynamic dimension size doesn't change the upper bound of the static
// shape.
let padded_v_five: f32[10] = set_dimension_size(v, five, /*dimension=*/0);
let padded_v_six: f32[10] = set_dimension_size(v, six, /*dimension=*/0);

// sum == 1 + 2 + 3 + 4 + 5
let sum:f32[] = reduce_sum(padded_v_five);
// product == 1 * 2 * 3 * 4 * 5
let product:f32[] = reduce_product(padded_v_five);

// Changing padding size will yield different result.
// sum == 1 + 2 + 3 + 4 + 5 + 6
let sum:f32[] = reduce_sum(padded_v_six);

GetTupleElement

Consulta también XlaBuilder::GetTupleElement.

Índices en una tupla con un valor constante de tiempo de compilación.

El valor debe ser una constante de tiempo de compilación para que la inferencia de forma pueda determinar el tipo del valor resultante.

Esto es análogo a std::get<int N>(t) en C++. Conceptualmente:

let v: f32[10] = f32[10]{0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
let s: s32 = 5;
let t: (f32[10], s32) = tuple(v, s);
let element_1: s32 = gettupleelement(t, 1);  // Inferred shape matches s32.

Consulta también tf.tuple.

Entrada

Consulta también XlaBuilder::Infeed.

Infeed(shape)

Argumento Tipo Semántica
shape Shape Es la forma de los datos leídos en la interfaz de Infeed. El campo de diseño de la forma se debe configurar para que coincida con el diseño de los datos enviados al dispositivo; de lo contrario, su comportamiento no está definido.

Lee un solo elemento de datos desde la interfaz de transmisión de entrada implícita del dispositivo, interpreta los datos como la forma determinada y su diseño, y muestra una XlaOp de los datos. Se permiten varias operaciones de entrada en un cálculo, pero debe haber un orden total entre ellas. Por ejemplo, dos entradas del siguiente código tienen un orden total, ya que hay una dependencia entre los bucles while.

result1 = while (condition, init = init_value) {
  Infeed(shape)
}

result2 = while (condition, init = result1) {
  Infeed(shape)
}

Las formas de tuplas anidadas no son compatibles. En el caso de una forma de tupla vacía, la operación de entrada es, de hecho, una no-op y continúa sin leer ningún dato de la entrada del dispositivo.

Pizca

Consulta también XlaBuilder::Iota.

Iota(shape, iota_dimension)

Compila un literal constante en el dispositivo en lugar de una transferencia de host potencialmente grande. Crea un array con una forma especificada y contiene valores que comienzan en cero y que aumentan en uno a lo largo de la dimensión especificada. Para los tipos de punto flotante, el array producido es equivalente a ConvertElementType(Iota(...)), en el que Iota es de tipo integral y la conversión es al tipo de punto flotante.

Argumentos Tipo Semántica
shape Shape Forma del array creado por Iota()
iota_dimension int64 Es la dimensión que se incrementa.

Por ejemplo, Iota(s32[4, 8], 0) muestra

  [[0, 0, 0, 0, 0, 0, 0, 0 ],
   [1, 1, 1, 1, 1, 1, 1, 1 ],
   [2, 2, 2, 2, 2, 2, 2, 2 ],
   [3, 3, 3, 3, 3, 3, 3, 3 ]]

Devoluciones por Iota(s32[4, 8], 1)

  [[0, 1, 2, 3, 4, 5, 6, 7 ],
   [0, 1, 2, 3, 4, 5, 6, 7 ],
   [0, 1, 2, 3, 4, 5, 6, 7 ],
   [0, 1, 2, 3, 4, 5, 6, 7 ]]

Asignar

Consulta también XlaBuilder::Map.

Map(operands..., computation)

Argumentos Tipo Semántica
operands secuencia de N XlaOp N matrices de tipos T0..T{N-1}
computation XlaComputation Cálculo de tipo T_0, T_1, .., T_{N + M -1} -> S con N parámetros de tipo T y M de tipo arbitrario
dimensions Array int64 array de dimensiones del mapa

Aplica una función escalar sobre los arrays operands dados, lo que produce un array con las mismas dimensiones en el que cada elemento es el resultado de la función asignada aplicada a los elementos correspondientes de los arrays de entrada.

La función asignada es un cálculo arbitrario con la restricción de que tiene N entradas del tipo escalar T y una sola salida con el tipo S. La salida tiene las mismas dimensiones que los operandos, excepto que el tipo de elemento T se reemplaza por S.

Por ejemplo: Map(op1, op2, op3, computation, par1) asigna elem_out <- computation(elem1, elem2, elem3, par1) a cada índice (multidimensional) en los arreglos de entrada para producir el arreglo de salida.

OptimizationBarrier

Bloquea cualquier pase de optimización para que no desplace los cálculos a través de la barrera.

Garantiza que todas las entradas se evalúen antes que los operadores que dependen de las salidas de la barrera.

Almohadilla

Consulta también XlaBuilder::Pad.

Pad(operand, padding_value, padding_config)

Argumentos Tipo Semántica
operand XlaOp array de tipo T
padding_value XlaOp escalar de tipo T para completar el padding agregado.
padding_config PaddingConfig la cantidad de padding en ambos bordes (bajo, alto) y entre los elementos de cada dimensión

Expande el array operand dado mediante padding alrededor de él y entre sus elementos con el padding_value especificado. padding_config especifica la cantidad de padding del borde y padding interior de cada dimensión.

PaddingConfig es un campo repetido de PaddingConfigDimension, que contiene tres campos para cada dimensión: edge_padding_low, edge_padding_high y interior_padding.

edge_padding_low y edge_padding_high especifican la cantidad de padding que se agrega en el extremo inferior (junto al índice 0) y en el extremo alto (junto al índice más alto) de cada dimensión, respectivamente. La cantidad de padding del borde puede ser negativa; el valor absoluto del padding negativo indica la cantidad de elementos que se quitarán de la dimensión especificada.

interior_padding especifica la cantidad de padding que se agrega entre dos elementos cualesquiera de cada dimensión. Este valor no puede ser negativo. El relleno interior se produce de forma lógica antes del relleno del borde, por lo que, en el caso del relleno del borde negativo, los elementos se quitan del operando con relleno interno.

Esta operación es una no-op si todos los pares de padding de bordes son (0, 0) y los valores de padding interior son todos 0. En la siguiente figura, se muestran ejemplos de diferentes valores edge_padding y interior_padding para un array bidimensional.

Recv

Consulta también XlaBuilder::Recv.

Recv(shape, channel_handle)

Argumentos Tipo Semántica
shape Shape forma de los datos para recibir
channel_handle ChannelHandle identificador único para cada par de envío/recepción

Recibe datos de la forma determinada de una instrucción Send en otro cálculo que comparte el mismo controlador de canal. Muestra una XlaOp para los datos recibidos.

La API de cliente de la operación Recv representa la comunicación síncrona. Sin embargo, la instrucción se divide de forma interna en 2 instrucciones HLO (Recv y RecvDone) para habilitar las transferencias de datos asíncronas. Consulta también HloInstruction::CreateRecv y HloInstruction::CreateRecvDone.

Recv(const Shape& shape, int64 channel_id)

Asigna los recursos necesarios para recibir datos de una instrucción Send con el mismo channel_id. Muestra un contexto para los recursos asignados, que se usa con una instrucción RecvDone siguiente a fin de esperar que se complete la transferencia de datos. El contexto es una tupla de {ángulo de recepción (forma), identificador de solicitud (U32)} y solo se puede usar mediante una instrucción RecvDone.

RecvDone(HloInstruction context)

En un contexto creado por una instrucción Recv, espera a que se complete la transferencia de datos y muestra los datos recibidos.

Reducir

Consulta también XlaBuilder::Reduce.

Aplica una función de reducción a uno o más arrays en paralelo.

Reduce(operands..., init_values..., computation, dimensions)

Argumentos Tipo Semántica
operands Secuencia de N XlaOp N arrays de tipos T_0, ..., T_{N-1}.
init_values Secuencia de N XlaOp N escalares de tipos T_0, ..., T_{N-1}.
computation XlaComputation el cálculo del tipo T_0, ..., T_{N-1}, T_0, ..., T_{N-1} -> Collate(T_0, ..., T_{N-1}).
dimensions Array int64 y sin ordenar de las dimensiones que se reducirán.

Aquí:

  • N debe ser mayor o igual que 1.
  • El cálculo debe ser asociativo “aproximadamente” (consulta a continuación).
  • Todos los arrays de entrada deben tener las mismas dimensiones.
  • Todos los valores iniciales deben formar una identidad inferior a computation.
  • Si es N = 1, Collate(T) es T.
  • Si es N > 1, Collate(T_0, ..., T_{N-1}) es una tupla de elementos N de tipo T.

Esta operación reduce una o más dimensiones de cada array de entrada en escalares. La clasificación de cada arreglo que se muestra es rank(operand) - len(dimensions). El resultado de la operación es Collate(Q_0, ..., Q_N), en el que Q_i es un array de tipo T_i, cuyas dimensiones se describen a continuación.

Diferentes backends pueden volver a asociar el cálculo de la reducción. Esto puede generar diferencias numéricas, ya que algunas funciones de reducción, como la suma, no son asociativas con los números de punto flotante. Sin embargo, si el rango de los datos es limitado, la adición de punto flotante está lo suficientemente cerca de ser asociativa para los usos más prácticos.

Ejemplos

Cuando se realiza una reducción en una dimensión en un solo array 1D con valores [10, 11, 12, 13], con la función de reducción f (esto es computation), eso se podría calcular como

f(10, f(11, f(12, f(init_value, 13)))

pero también existen muchas otras posibilidades, p.ej.

f(init_value, f(f(10, f(init_value, 11)), f(f(init_value, 12), f(init_value, 13))))

El siguiente es un ejemplo aproximado de un seudocódigo de cómo se puede implementar la reducción, con la sumatoria como el cálculo de la reducción con un valor inicial de 0.

result_shape <- remove all dims in dimensions from operand_shape

# Iterate over all elements in result_shape. The number of r's here is equal
# to the rank of the result
for r0 in range(result_shape[0]), r1 in range(result_shape[1]), ...:
  # Initialize this result element
  result[r0, r1...] <- 0

  # Iterate over all the reduction dimensions
  for d0 in range(dimensions[0]), d1 in range(dimensions[1]), ...:
    # Increment the result element with the value of the operand's element.
    # The index of the operand's element is constructed from all ri's and di's
    # in the right order (by construction ri's and di's together index over the
    # whole operand shape).
    result[r0, r1...] += operand[ri... di]

Este es un ejemplo de cómo reducir un array 2D (matriz). La forma tiene rango 2, dimensión 0 de tamaño 2 y dimensión 1 de tamaño 3:

Resultados de reducir las dimensiones 0 o 1 con una función "add":

Ten en cuenta que ambos resultados de reducción son arrays 1D. En el diagrama, se muestra uno como columna y otro como fila solo para mayor conveniencia visual.

Para un ejemplo más complejo, aquí hay un array 3D. Su clasificación es 3, dimensión 0 de tamaño 4, dimensión 1 de tamaño 2 y dimensión 2 de tamaño 3. Para simplificar, los valores 1 a 6 se replican en la dimensión 0.

Al igual que en el ejemplo en 2D, podemos reducir solo una dimensión. Si reducimos la dimensión 0, por ejemplo, obtenemos un array de rango 2 en el que todos los valores de la dimensión 0 se pliegan en un escalar:

|  4   8  12 |
| 16  20  24 |

Si reducimos la dimensión 2, también obtenemos un array de rango 2 en el que todos los valores de la dimensión 2 se plegaron en un escalar:

| 6  15 |
| 6  15 |
| 6  15 |
| 6  15 |

Ten en cuenta que el orden relativo entre las dimensiones restantes de la entrada se conserva en la salida, pero es posible que a algunas dimensiones se les asignen números nuevos (debido a que cambia la clasificación).

También podemos reducir varias dimensiones. De reducción de las dimensiones 0 y 1, se produce un array 1D [20, 28, 36].

Si reduces el array 3D en todas sus dimensiones, se produce el 84 escalar.

Reducción variable

Cuando es N > 1, la aplicación de la función reducir es un poco más compleja, ya que se aplica simultáneamente a todas las entradas. Los operandos se suministran al cálculo en el siguiente orden:

  • Ejecutando valor reducido para el primer operando
  • ...
  • Ejecutando valor reducido para el operando enésimo
  • Valor de entrada para el primer operando
  • ...
  • Valor de entrada para el operando enésimo

Por ejemplo, considera la siguiente función de reducción, que se puede usar para calcular el máximo y el argmax de un array de 1D en paralelo:

f: (Float, Int, Float, Int) -> Float, Int
f(max, argmax, value, index):
  if value >= max:
    return (value, index)
  else:
    return (max, argmax)

Para los arrays de entrada 1D V = Float[N], K = Int[N] y los valores init I_V = Float, I_K = Int, el resultado f_(N-1) de la reducción en la única dimensión de entrada es equivalente a la siguiente aplicación recursiva:

f_0 = f(I_V, I_K, V_0, K_0)
f_1 = f(f_0.first, f_0.second, V_1, K_1)
...
f_(N-1) = f(f_(N-2).first, f_(N-2).second, V_(N-1), K_(N-1))

Si aplicas esta reducción a un array de valores y un array de índices secuenciales (es decir, iota), se iterarán en conjunto los arrays y se mostrará una tupla que contenga el valor máximo y el índice coincidente.

ReducePrecision

Consulta también XlaBuilder::ReducePrecision.

Modela el efecto de convertir valores de punto flotante a un formato de menor precisión (como IEEE-FP16) y de vuelta al formato original. La cantidad de bits de exponentes y mantisas en el formato de menor precisión se puede especificar de manera arbitraria, aunque es posible que no todos los tamaños de bits sean compatibles con todas las implementaciones de hardware.

ReducePrecision(operand, mantissa_bits, exponent_bits)

Argumentos Tipo Semántica
operand XlaOp array de tipo de punto flotante T.
exponent_bits int32 cantidad de bits exponentes en un formato de menor precisión
mantissa_bits int32 cantidad de bits de mantisa en un formato de menor precisión

El resultado es un array de tipo T. Los valores de entrada se redondean al valor más cercano representable con la cantidad determinada de bits de mantisa (mediante la semántica de "empates a pares"), y cualquier valor que exceda el rango especificado por la cantidad de bits exponentes se restringen al infinito positivo o negativo. Se conservan los valores NaN, aunque se pueden convertir en valores NaN canónicos.

El formato de menor precisión debe tener al menos un bit del exponente (para distinguir un valor cero de un valor infinito, ya que ambos tienen una mantisa cero) y debe tener un número no negativo de bits de mantisa. La cantidad de bits de exponente o mantisa puede exceder el valor correspondiente para el tipo T. La parte correspondiente de la conversión es simplemente una no-op.

ReduceScatter

Consulta también XlaBuilder::ReduceScatter.

ReduceScatter es una operación colectiva que realiza un AllReduce y, luego, dispersa el resultado dividiéndolo en bloques shard_count a lo largo de scatter_dimension y la réplica i del grupo de réplicas recibe el fragmento ith.

ReduceScatter(operand, computation, scatter_dim, shard_count, replica_group_ids, channel_id)

Argumentos Tipo Semántica
operand XlaOp Es un array o una tupla no vacía de arrays que se reducirán en las réplicas.
computation XlaComputation Cálculo de reducción
scatter_dimension int64 Dimensión que se debe dispersar.
shard_count int64 Cantidad de bloques para dividir scatter_dimension
replica_groups vector de vectores de int64 Grupos entre los que se realizan las reducciones
channel_id int64 opcional ID de canal opcional para la comunicación entre módulos
  • Cuando operand es una tupla de arrays, se realiza la reducción de dispersión en cada elemento de la tupla.
  • replica_groups es una lista de grupos de réplicas entre los que se realiza la reducción (el ID de réplica de la réplica actual se puede recuperar con ReplicaId). El orden de las réplicas en cada grupo determina el orden en el que se dispersará el resultado de All-reduce. replica_groups debe estar vacío (en ese caso, todas las réplicas pertenecen a un solo grupo) o contener la misma cantidad de elementos que la cantidad de réplicas. Cuando hay más de un grupo de réplicas, todos deben tener el mismo tamaño. Por ejemplo, replica_groups = {0, 2}, {1, 3} realiza una reducción entre las réplicas 0 y 2, y 1 y 3, y, luego, dispersa el resultado.
  • shard_count es el tamaño de cada grupo de réplicas. Esto lo necesitamos en los casos en que los objetos replica_groups están vacíos. Si replica_groups no está vacío, shard_count debe ser igual al tamaño de cada grupo de réplicas.
  • channel_id se usa para la comunicación entre módulos: solo las operaciones reduce-scatter con el mismo channel_id pueden comunicarse entre sí.

La forma de salida es la de entrada, con scatter_dimension hecho shard_count veces más pequeño. Por ejemplo, si hay dos réplicas y el operando tiene el valor [1.0, 2.25] y [3.0, 5.25] respectivamente en ambas réplicas, el valor de salida de esta operación en la que scatter_dim es 0 será [4.0] para la primera réplica y [7.5] para la segunda.

ReduceWindow

Consulta también XlaBuilder::ReduceWindow.

Aplica una función de reducción a todos los elementos en cada ventana de una secuencia de N arrays multidimensionales, lo que produce una única o una tupla de N arrays multidimensionales como resultado. Cada arreglo de salida tiene la misma cantidad de elementos que la cantidad de posiciones válidas de la ventana. Una capa de reducción se puede expresar como ReduceWindow. Al igual que Reduce, al objeto computation aplicado siempre se le pasa el init_values en el lado izquierdo.

ReduceWindow(operands..., init_values..., computation, window_dimensions, window_strides, padding)

Argumentos Tipo Semántica
operands N XlaOps Es una secuencia de N arrays multidimensionales de tipos T_0,..., T_{N-1}, cada uno de los cuales representa el área base en la que se coloca la ventana.
init_values N XlaOps Los N valores iniciales para la reducción, uno para cada uno de los N operandos. Consulta Reducir para obtener más detalles.
computation XlaComputation Función de reducción de tipo T_0, ..., T_{N-1}, T_0, ..., T_{N-1} -> Collate(T_0, ..., T_{N-1}), para aplicar a elementos en cada ventana de todos los operandos de entrada.
window_dimensions ArraySlice<int64> array de números enteros para los valores de las dimensiones de ventana
window_strides ArraySlice<int64> Array de números enteros para los valores de segmento de ventana
base_dilations ArraySlice<int64> array de números enteros para valores de dilatación base
window_dilations ArraySlice<int64> array de números enteros para valores de dilatación de ventanas
padding Padding el tipo de padding para la ventana (Padding::kSame, que se rellena para tener la misma forma de salida que la entrada si el stride es 1, o Padding::kValid, que no usa relleno y “detiene” la ventana cuando ya no encaja)

Aquí:

  • N debe ser mayor o igual que 1.
  • Todos los arrays de entrada deben tener las mismas dimensiones.
  • Si es N = 1, Collate(T) es T.
  • Si es N > 1, Collate(T_0, ..., T_{N-1}) es una tupla de elementos N de tipo (T0,...T{N-1}).

Debajo del código y la figura, se muestra un ejemplo del uso de ReduceWindow. La entrada es una matriz de tamaño [4x6], y tanto window_dimensions como window_stride_dimensions son [2x3].

// Create a computation for the reduction (maximum).
XlaComputation max;
{
  XlaBuilder builder(client_, "max");
  auto y = builder.Parameter(0, ShapeUtil::MakeShape(F32, {}), "y");
  auto x = builder.Parameter(1, ShapeUtil::MakeShape(F32, {}), "x");
  builder.Max(y, x);
  max = builder.Build().value();
}

// Create a ReduceWindow computation with the max reduction computation.
XlaBuilder builder(client_, "reduce_window_2x3");
auto shape = ShapeUtil::MakeShape(F32, {4, 6});
auto input = builder.Parameter(0, shape, "input");
builder.ReduceWindow(
    input,
    /*init_val=*/builder.ConstantLiteral(LiteralUtil::MinValue(F32)),
    *max,
    /*window_dimensions=*/{2, 3},
    /*window_stride_dimensions=*/{2, 3},
    Padding::kValid);

El segmento de 1 en una dimensión especifica que la posición de una ventana en la dimensión se encuentra a 1 elemento de distancia de la ventana adyacente. Para especificar que ninguna ventana se superpone entre sí, window_stride_dimensions debe ser igual a window_dimensions. En la siguiente figura, se ilustra el uso de dos valores de segmento diferentes. El relleno se aplica a cada dimensión de la entrada y los cálculos son los mismos que si la entrada hubiera llegado con las dimensiones que tenía después del relleno.

Para un ejemplo de padding no trivial, considera calcular el mínimo de ventana reducida (el valor inicial es MAX_FLOAT) con la dimensión 3 y el segmento 2 sobre el array de entrada [10000, 1000, 100, 10, 1]. El padding kValid calcula los mínimos en dos ventanas válidas: [10000, 1000, 100] y [100, 10, 1], lo que genera el resultado [100, 1]. Con el padding kSame, primero se rellena el array para que la forma después de la ventana de reducción sea igual que la entrada para el segmento uno. Para ello, agrega elementos iniciales en ambos lados y obtén [MAX_VALUE, 10000, 1000, 100, 10, 1, MAX_VALUE]. La ejecución de la función "reduce-window" en el array con padding opera en tres ventanas: [MAX_VALUE, 10000, 1000], [1000, 100, 10], [10, 1, MAX_VALUE] y produce [1000, 10, 1].

El orden de evaluación de la función de reducción es arbitrario y puede ser no determinista. Por lo tanto, la función de reducción no debe ser demasiado sensible a la reasociación. Consulta el debate sobre la asociatividad en el contexto de Reduce para obtener más detalles.

ReplicaId

Consulta también XlaBuilder::ReplicaId.

Muestra el ID único (escalar U32) de la réplica.

ReplicaId()

El ID único de cada réplica es un número entero sin firmar en el intervalo [0, N), en el que N es la cantidad de réplicas. Dado que todas las réplicas ejecutan el mismo programa, una llamada ReplicaId() en el programa mostrará un valor diferente en cada réplica.

Reshape

Consulta también XlaBuilder::Reshape y la operación Collapse.

Cambia la forma de las dimensiones de un array a una configuración nueva.

Reshape(operand, new_sizes) Reshape(operand, dimensions, new_sizes)

Argumentos Tipo Semántica
operand XlaOp array de tipo T
dimensions Vector de int64 orden en el que se contraen las dimensiones
new_sizes Vector de int64 vector de tamaños de nuevas dimensiones

Conceptualmente, primero, redimensionar compacta un array en un vector unidimensional de valores de datos y, luego, define mejor este vector en una nueva forma. Los argumentos de entrada son un arreglo arbitrario de tipo T, un vector de constante de tiempo de compilación de índices de dimensión y un vector de constante de tiempo de compilación de tamaños de dimensión para el resultado. Los valores en el vector dimension, si se proporcionan, deben ser una permutación de todas las dimensiones de T; el valor predeterminado si no se proporciona es {0, ..., rank - 1}. El orden de las dimensiones en dimensions va desde la dimensión más lenta que varía (la más importante) hasta la que varía más rápida (la más secundaria) en el anidamiento de bucles que contrae el array de entrada en una sola dimensión. El vector new_sizes determina el tamaño del array de salida. El valor del índice 0 de new_sizes es el tamaño de la dimensión 0, el valor del índice 1 es el tamaño de la dimensión 1, y así sucesivamente. El producto de las dimensiones new_size debe ser igual al producto de los tamaños de dimensión del operando. Cuando defines mejor el array contraído en el array multidimensional definido por new_sizes, las dimensiones en new_sizes se ordenan desde la variación más lenta (la más importante) hasta la más rápida (la más secundaria).

Por ejemplo, supongamos que v es un array de 24 elementos:

let v = f32[4x2x3] { { {10, 11, 12}, {15, 16, 17} },
                    { {20, 21, 22}, {25, 26, 27} },
                    { {30, 31, 32}, {35, 36, 37} },
                    { {40, 41, 42}, {45, 46, 47} } };

In-order collapse:
let v012_24 = Reshape(v, {0,1,2}, {24});
then v012_24 == f32[24] {10, 11, 12, 15, 16, 17, 20, 21, 22, 25, 26, 27,
                         30, 31, 32, 35, 36, 37, 40, 41, 42, 45, 46, 47};

let v012_83 = Reshape(v, {0,1,2}, {8,3});
then v012_83 == f32[8x3] { {10, 11, 12}, {15, 16, 17},
                          {20, 21, 22}, {25, 26, 27},
                          {30, 31, 32}, {35, 36, 37},
                          {40, 41, 42}, {45, 46, 47} };

Out-of-order collapse:
let v021_24 = Reshape(v, {1,2,0}, {24});
then v012_24 == f32[24]  {10, 20, 30, 40, 11, 21, 31, 41, 12, 22, 32, 42,
                          15, 25, 35, 45, 16, 26, 36, 46, 17, 27, 37, 47};

let v021_83 = Reshape(v, {1,2,0}, {8,3});
then v021_83 == f32[8x3] { {10, 20, 30}, {40, 11, 21},
                          {31, 41, 12}, {22, 32, 42},
                          {15, 25, 35}, {45, 16, 26},
                          {36, 46, 17}, {27, 37, 47} };


let v021_262 = Reshape(v, {1,2,0}, {2,6,2});
then v021_262 == f32[2x6x2] { { {10, 20}, {30, 40},
                              {11, 21}, {31, 41},
                              {12, 22}, {32, 42} },
                             { {15, 25}, {35, 45},
                              {16, 26}, {36, 46},
                              {17, 27}, {37, 47} } };

Como caso especial, remodelar puede transformar un array de un solo elemento en un escalar y viceversa. Por ejemplo,

Reshape(f32[1x1] { {5} }, {0,1}, {}) == 5;
Reshape(5, {}, {1,1}) == f32[1x1] { {5} };

Rev (reversión)

Consulta también XlaBuilder::Rev.

Rev(operand, dimensions)

Argumentos Tipo Semántica
operand XlaOp array de tipo T
dimensions ArraySlice<int64> dimensiones para revertir

Invierte el orden de los elementos en el array operand a lo largo del dimensions especificado, lo que genera un array de salida de la misma forma. Cada elemento del array de operandos en un índice multidimensional se almacena en el array de salida en un índice transformado. Para transformar el índice multidimensional, se invierte el índice de cada dimensión que se revertirá (es decir, si una dimensión de tamaño N es una de las dimensiones invertidas, su índice i se transforma en N - 1 - i).

Un uso de la operación Rev es revertir el array de pesos de convolución junto con las dos dimensiones de ventana durante el cálculo de gradientes en redes neuronales.

RngNormal

Consulta también XlaBuilder::RngNormal.

Construye un resultado de una forma dada con números al azar generados siguiendo la \(N(\mu, \sigma)\) distribución normal. Los parámetros \(\mu\) , \(\sigma\)y la forma de salida deben tener un tipo de elemento de punto flotante. Además, los parámetros deben tener valores escalares.

RngNormal(mu, sigma, shape)

Argumentos Tipo Semántica
mu XlaOp Escalar de tipo T que especifica la media de los números generados
sigma XlaOp Escalar de tipo T que especifica la desviación estándar del modelo generado
shape Shape Forma de salida de tipo T

RngUniform

Consulta también XlaBuilder::RngUniform.

Construye un resultado de una forma determinada con números aleatorios generados a partir de la distribución uniforme durante el intervalo \([a,b)\). Los parámetros y el tipo de elemento de salida deben ser de tipo booleano, integral o de punto flotante, y los tipos deben ser coherentes. En la actualidad, los backends de CPU y GPU solo admiten F64, F32, F16, BF16, S64, U64, S32 y U32. Además, los parámetros deben tener valores escalares. Si \(b <= a\) el resultado está definido por la implementación.

RngUniform(a, b, shape)

Argumentos Tipo Semántica
a XlaOp Escalar de tipo T que especifica el límite de intervalo más bajo
b XlaOp Escalar de tipo T que especifica el límite superior del intervalo
shape Shape Forma de salida de tipo T

RngBitGenerator

Genera un resultado con una forma determinada rellena de bits aleatorios uniformes mediante el algoritmo especificado (o el valor predeterminado del backend) y muestra un estado actualizado (con la misma forma que el estado inicial) y los datos aleatorios generados.

El estado inicial es el estado inicial de la generación actual de números aleatorios. Este, la forma requerida y los valores válidos dependen del algoritmo utilizado.

Se garantiza que el resultado será una función determinista del estado inicial, pero no se garantiza que lo sea entre los backends y las diferentes versiones del compilador.

RngBitGenerator(algorithm, key, shape)

Argumentos Tipo Semántica
algorithm RandomAlgorithm Se usará el algoritmo PRNG.
initial_state XlaOp Estado inicial del algoritmo PRNG
shape Shape Forma de la salida para los datos generados.

Valores disponibles para algorithm:

Dispersión

La operación de dispersión de XLA genera una secuencia de resultados que son valores del arreglo de entrada operands, con varias porciones (en los índices especificados por scatter_indices) actualizadas con la secuencia de valores en updates mediante update_computation.

Consulta también XlaBuilder::Scatter.

scatter(operands..., scatter_indices, updates..., update_computation, index_vector_dim, update_window_dims, inserted_window_dims, scatter_dims_to_operand_dims)

Argumentos Tipo Semántica
operands Secuencia de N XlaOp N arrays de tipos T_0, ..., T_N en los que se dispersará.
scatter_indices XlaOp Arreglo que contiene los índices iniciales de las porciones en las que se deben dispersar.
updates Secuencia de N XlaOp N arrays de tipos T_0, ..., T_N. updates[i] contiene los valores que deben usarse para la dispersión operands[i].
update_computation XlaComputation Cálculo que se usa para combinar los valores existentes en el array de entrada y las actualizaciones durante la dispersión. Este cálculo debe ser del tipo T_0, ..., T_N, T_0, ..., T_N -> Collate(T_0, ..., T_N).
index_vector_dim int64 Es la dimensión en scatter_indices que contiene los índices iniciales.
update_window_dims ArraySlice<int64> Es el conjunto de dimensiones en forma updates que son dimensiones de ventana.
inserted_window_dims ArraySlice<int64> Es el conjunto de dimensiones de la ventana que se deben insertar con la forma updates.
scatter_dims_to_operand_dims ArraySlice<int64> Se asignan dimensiones de los índices de dispersión al espacio de índices de operando. Este array se interpreta como la asignación de i a scatter_dims_to_operand_dims[i] . Tienen que ser uno a uno y totales.
indices_are_sorted bool Indica si se garantiza que los índices estén ordenados por el emisor.
unique_indices bool Indica si se garantiza que los índices sean únicos para el emisor.

Aquí:

  • N debe ser mayor o igual que 1.
  • operands[0], ... y operands[N-1] deben tener las mismas dimensiones.
  • updates[0], ... y updates[N-1] deben tener las mismas dimensiones.
  • Si es N = 1, Collate(T) es T.
  • Si es N > 1, Collate(T_0, ..., T_N) es una tupla de elementos N de tipo T.

Si index_vector_dim es igual a scatter_indices.rank, de manera implícita consideramos que scatter_indices tiene una dimensión 1 final.

Definimos update_scatter_dims de tipo ArraySlice<int64> como el conjunto de dimensiones en la forma updates que no están en update_window_dims, en orden ascendente.

Los argumentos de dispersión deben seguir estas restricciones:

  • Cada array updates debe tener la clasificación update_window_dims.size + scatter_indices.rank - 1.

  • Los límites de la dimensión i en cada array updates deben cumplir con lo siguiente:

    • Si i está presente en update_window_dims (es decir, igual a update_window_dims[k] para alguna k), el límite de la dimensión i en updates no debe exceder el límite correspondiente de operand después de contemplar inserted_window_dims (es decir, adjusted_window_bounds[k], donde adjusted_window_bounds contiene los límites de operand sin los límites de los índices inserted_window_dims).
    • Si i está presente en update_scatter_dims (es decir, igual a update_scatter_dims[k] para algunos k), el límite de la dimensión i en updates debe ser igual al límite correspondiente de scatter_indices y se omite index_vector_dim (es decir, scatter_indices.shape.dims[k], si k < index_vector_dim y scatter_indices.shape.dims[k+1]).
  • update_window_dims debe estar en orden ascendente, no debe tener ningún número de dimensión que se repita y estar en el rango [0, updates.rank).

  • inserted_window_dims debe estar en orden ascendente, no debe tener ningún número de dimensión que se repita y estar en el rango [0, operand.rank).

  • operand.rank debe ser igual a la suma de update_window_dims.size y inserted_window_dims.size.

  • scatter_dims_to_operand_dims.size debe ser igual a scatter_indices.shape.dims[index_vector_dim] y sus valores deben estar en el rango [0, operand.rank).

Para un índice determinado U en cada array updates, el índice I correspondiente en el array operands correspondiente al que se debe aplicar esta actualización se calcula de la siguiente manera:

  1. Deja que G = { U[k] para k en update_scatter_dims }. Usa G para buscar un vector de índice S en el array scatter_indices de modo que S[i] = scatter_indices[Combine(G, i)] donde Combine(A, b) inserte b en las posiciones index_vector_dim en A.
  2. Crea un índice Sin en operand con S mediante la dispersión S con el mapa scatter_dims_to_operand_dims. Más formalmente:
    1. Sin[scatter_dims_to_operand_dims[k]] = S[k] si k < scatter_dims_to_operand_dims.size.
    2. Sin[_] = 0 de lo contrario.
  3. Crea un índice Win en cada array operands. Para ello, dispersa los índices en update_window_dims en U de acuerdo con inserted_window_dims. Más formalmente:
    1. Win[window_dims_to_operand_dims(k)] = U[k] si k está en update_window_dims, donde window_dims_to_operand_dims es la función monótona con el dominio [0, update_window_dims.size) y el rango [0, operand.rank) \ inserted_window_dims. (Por ejemplo, si update_window_dims.size es 4, operand.rank es 6, y inserted_window_dims es {0, 2}, window_dims_to_operand_dims es {01, 13, 24, 35}).
    2. Win[_] = 0 de lo contrario.
  4. I es Win + Sin, donde + es la suma de los elementos.

En resumen, la operación de dispersión puede definirse de la siguiente manera.

  • Inicializa output con operands, es decir, para todos los índices J, para todos los índices O del array operands[J]:
    output[J][O] = operands[J][O]
  • Para cada índice U del array updates[J] y el índice correspondiente O en el array operand[J], si O es un índice válido para output:
    (output[0][O], ..., output[N-1][O]) =update_computation(output[0][O], ..., ,output[N-1][O],updates[0][U], ...,updates[N-1][U])

El orden en el que se aplican las actualizaciones no es determinista. Por lo tanto, cuando varios índices de updates hacen referencia al mismo índice en operands, el valor correspondiente en output no será determinista.

Ten en cuenta que el primer parámetro que se pasa a update_computation siempre será el valor actual del array output y el segundo siempre será el valor del array updates. Esto es importante en particular para los casos en los que update_computation no es conmutativo.

Si indices_are_sorted se configura como verdadero, XLA puede suponer que el usuario ordena start_indices (en orden ascendente start_index_map). Si no lo están, entonces se define la semántica de la implementación.

Si unique_indices se establece como verdadero, XLA puede suponer que todos los elementos dispersos son únicos. Por lo tanto, XLA podría usar operaciones no atómicas. Si unique_indices se configura como verdadero y los índices que se dispersan no son únicos, se define la implementación de la semántica.

informalmente, la op de dispersión puede verse como una inversión de la op de recopilación, es decir, la op de dispersión actualiza los elementos en la entrada que son extraídos por la op de recopilación correspondiente.

Para obtener una descripción informal detallada y ejemplos, consulta la sección "Descripción informal" en Gather.

Seleccionar

Consulta también XlaBuilder::Select.

Construye un arreglo de salida a partir de elementos de dos arreglos de entrada, en función de los valores de un arreglo de predicado.

Select(pred, on_true, on_false)

Argumentos Tipo Semántica
pred XlaOp array de tipo PRED
on_true XlaOp array de tipo T
on_false XlaOp array de tipo T

Los arrays on_true y on_false deben tener la misma forma. Esta también es la forma del array de salida. El array pred debe tener la misma dimensionalidad que on_true y on_false, con el tipo de elemento PRED.

Para cada elemento P de pred, el elemento correspondiente del array de salida se toma de on_true si el valor de P es true y de on_false si el valor de P es false. Como una forma restringida de transmisión, pred puede ser un escalar de tipo PRED. En este caso, el array de salida se toma por completo de on_true si pred es true y de on_false si pred es false.

Ejemplo con pred no escalar:

let pred: PRED[4] = {true, false, false, true};
let v1: s32[4] = {1, 2, 3, 4};
let v2: s32[4] = {100, 200, 300, 400};
==>
Select(pred, v1, v2) = s32[4]{1, 200, 300, 4};

Ejemplo con pred escalar:

let pred: PRED = true;
let v1: s32[4] = {1, 2, 3, 4};
let v2: s32[4] = {100, 200, 300, 400};
==>
Select(pred, v1, v2) = s32[4]{1, 2, 3, 4};

Se admiten las selecciones entre tuplas. Para este propósito, se considera que las tuplas son tipos escalares. Si on_true y on_false son tuplas (que deben tener la misma forma), pred debe ser un escalar de tipo PRED.

SelectAndScatter

Consulta también XlaBuilder::SelectAndScatter.

Esta operación puede considerarse como una operación compuesta que primero calcula ReduceWindow en el array operand para seleccionar un elemento de cada ventana y, luego, dispersa el array source en los índices de los elementos seleccionados para construir un array de salida con la misma forma que el array de operando. La función binaria select se usa para seleccionar un elemento de cada ventana aplicándolo en cada ventana, y se llama con la propiedad de que el vector de índice del primer parámetro es lexicográficamente menor que el vector de índice del segundo parámetro. La función select muestra true si se selecciona el primer parámetro y false si se selecciona el segundo, y la función debe contener transitividad (es decir, si select(a, b) y select(b, c) son true, entonces select(a, c) también es true) para que el elemento seleccionado no dependa del orden de los elementos recorridos en una ventana determinada.

La función scatter se aplica en cada índice seleccionado en el array de salida. Requiere dos parámetros escalares:

  1. Valor actual en el índice seleccionado en el array de salida
  2. El valor de dispersión de source que se aplica al índice seleccionado

Combina los dos parámetros y muestra un valor escalar que se usa para actualizar el valor en el índice seleccionado en el arreglo de salida. Inicialmente, todos los índices del arreglo de salida se configuran en init_value.

El array de salida tiene la misma forma que el array operand, y el array source debe tener la misma forma que el resultado de aplicar una operación ReduceWindow en el array operand. Se puede usar SelectAndScatter para propagar de forma inversa los valores de gradiente de una capa de reducción en una red neuronal.

SelectAndScatter(operand, select, window_dimensions, window_strides, padding, source, init_value, scatter)

Argumentos Tipo Semántica
operand XlaOp Array de tipo T sobre el que se deslizan las ventanas
select XlaComputation Cálculo binario de tipo T, T -> PRED, para aplicar a todos los elementos en cada ventana. Muestra true si se selecciona el primer parámetro y false si se selecciona el segundo parámetro.
window_dimensions ArraySlice<int64> array de números enteros para los valores de las dimensiones de ventana
window_strides ArraySlice<int64> Array de números enteros para los valores de segmento de ventana
padding Padding Tipo de padding para la ventana (Padding::kSame o Padding::kValid)
source XlaOp un array de tipo T con los valores que se van a dispersar.
init_value XlaOp valor escalar de tipo T para el valor inicial del array de salida
scatter XlaComputation Cálculo binario de tipo T, T -> T, para aplicar cada elemento de fuente de dispersión con su elemento de destino

En la siguiente figura, se muestran ejemplos del uso de SelectAndScatter con la función select que calcula el valor máximo entre sus parámetros. Ten en cuenta que cuando las ventanas se superponen, como en la figura (2) de abajo, diferentes ventanas pueden seleccionar un índice del array operand varias veces. En la figura, las dos ventanas superiores (azul y rojo) seleccionan el elemento de valor 9 y la función scatter de suma binaria produce el elemento de resultado del valor 8 (2 + 6).

El orden de evaluación de la función scatter es arbitrario y puede ser no determinista. Por lo tanto, la función scatter no debe ser muy sensible a la reasociación. Consulta el debate sobre la asociatividad en el contexto de Reduce para obtener más detalles.

Enviar

Consulta también XlaBuilder::Send.

Send(operand, channel_handle)

Argumentos Tipo Semántica
operand XlaOp datos para enviar (matriz de tipo T)
channel_handle ChannelHandle identificador único para cada par de envío/recepción

Envía los datos del operando dados a una instrucción Recv en otro cálculo que comparte el mismo controlador de canal. No muestra ningún dato.

Al igual que la operación Recv, la API de cliente de la operación Send representa la comunicación síncrona y se divide de manera interna en 2 instrucciones HLO (Send y SendDone) para habilitar las transferencias de datos asíncronas. Consulta también HloInstruction::CreateSend y HloInstruction::CreateSendDone.

Send(HloInstruction operand, int64 channel_id)

Inicia una transferencia asíncrona del operando a los recursos asignados por la instrucción Recv con el mismo ID de canal. Muestra un contexto, que se usa en una instrucción SendDone siguiente para esperar que se complete la transferencia de datos. El contexto es una tupla de {operando (forma), identificador de solicitud (U32)} y solo se puede usar con una instrucción SendDone.

SendDone(HloInstruction context)

En un contexto creado por una instrucción Send, espera a que se complete la transferencia de datos. Esta instrucción no muestra ningún dato.

Programación de instrucciones del canal

El orden de ejecución de las 4 instrucciones para cada canal (Recv, RecvDone, Send, SendDone) es el siguiente.

  • Recv ocurre antes del Send
  • Send ocurre antes del RecvDone
  • Recv ocurre antes del RecvDone
  • Send ocurre antes del SendDone

Cuando los compiladores de backend generan una programación lineal para cada procesamiento que se comunica a través de las instrucciones del canal, no debe haber ciclos a través de los cálculos. Por ejemplo, los siguientes programas generan interbloqueos.

Porción

Consulta también XlaBuilder::Slice.

La segmentación extrae un subarray del array de entrada. El subarreglo es del mismo rango que la entrada y contiene los valores dentro de un cuadro delimitador dentro del array de entrada, en el que las dimensiones y los índices del cuadro de límite se proporcionan como argumentos a la operación de división.

Slice(operand, start_indices, limit_indices, strides)

Argumentos Tipo Semántica
operand XlaOp Array N dimensional de tipo T
start_indices ArraySlice<int64> Lista de N números enteros que contienen los índices iniciales de la porción para cada dimensión. Los valores deben ser mayores o iguales que cero.
limit_indices ArraySlice<int64> Lista de N números enteros que contienen los índices finales (exclusivos) para la porción de cada dimensión. Cada valor debe ser mayor o igual que el valor de start_indices respectivo de la dimensión, así como menor o igual que el tamaño de la dimensión.
strides ArraySlice<int64> Lista de N números enteros que deciden el segmento de entrada de la porción. La porción elige cada elemento strides[d] en la dimensión d.

Ejemplo unidimensional:

let a = {0.0, 1.0, 2.0, 3.0, 4.0}
Slice(a, {2}, {4}) produces:
  {2.0, 3.0}

Ejemplo bidimensional:

let b =
 { {0.0,  1.0,  2.0},
   {3.0,  4.0,  5.0},
   {6.0,  7.0,  8.0},
   {9.0, 10.0, 11.0} }

Slice(b, {2, 1}, {4, 3}) produces:
  { { 7.0,  8.0},
    {10.0, 11.0} }

Ordenar

Consulta también XlaBuilder::Sort.

Sort(operands, comparator, dimension, is_stable)

Argumentos Tipo Semántica
operands ArraySlice<XlaOp> Los operandos que se van a ordenar.
comparator XlaComputation El cálculo del comparador que se usará.
dimension int64 La dimensión en la que se ordenará.
is_stable bool Si se debe usar la ordenación estable.

Si solo se proporciona un operando:

  • Si el operando es un tensor de rango-1 (un array), el resultado es un array ordenado. Si deseas ordenar el array en orden ascendente, el comparador debe realizar una comparación de menor que. De manera formal, después de ordenar el array, se mantiene en todas las posiciones de índice i, j con i < j que comparator(value[i], value[j]) = comparator(value[j], value[i]) = false o comparator(value[i], value[j]) = true.

  • Si el operando tiene una clasificación más alta, se ordena según la dimensión proporcionada. Por ejemplo, para un tensor de rango-2 (una matriz), un valor de dimensión de 0 ordenará cada columna de forma independiente, y un valor de dimensión de 1 ordenará cada fila de forma independiente. Si no se proporciona un número de dimensión, se elige la última de forma predeterminada. Para la dimensión que está ordenada, se aplica el mismo orden de clasificación que en el caso de rango-1.

Si se proporcionan operandos n > 1:

  • Todos los operandos n deben ser tensores con las mismas dimensiones. Los tipos de elementos de los tensores pueden ser diferentes.

  • Todos los operandos se ordenan juntos, no de manera individual. Conceptualmente, los operandos se tratan como una tupla. Cuando se verifica si los elementos de cada operando en las posiciones de índice i y j deben intercambiarse, se llama al comparador con los parámetros escalares 2 * n, en los que el parámetro 2 * k corresponde al valor en la posición i del operando k-th y el parámetro 2 * k + 1 corresponde al valor en la posición j del operando k-th. Por lo general, el comparador compararía los parámetros 2 * k y 2 * k + 1 entre sí y, posiblemente, usaría otros pares de parámetros como rompientes de empate.

  • El resultado es una tupla que consta de los operandos en orden (en la dimensión proporcionada, como se muestra arriba). El operando i-th de la tupla corresponde al operando i-th de Sort.

Por ejemplo, si hay tres operandos operand0 = [3, 1], operand1 = [42, 50], operand2 = [-3.0, 1.1], y el comparador solo compara los valores de operand0 con menor que, el resultado del orden es la tupla ([1, 3], [50, 42], [1.1, -3.0]).

Si is_stable se configura como verdadero, se garantiza que el orden sea estable, es decir, si hay elementos que el comparador considera iguales, se conserva el orden relativo de los valores iguales. Dos elementos e1 y e2 son iguales solo si comparator(e1, e2) = comparator(e2, e1) = false. Según la configuración predeterminada, is_stable se configura como falso.

Transposición

Consulta también la operación tf.reshape.

Transpose(operand)

Argumentos Tipo Semántica
operand XlaOp El operando que se va a transponer.
permutation ArraySlice<int64> Cómo permutar las dimensiones

Permuta las dimensiones del operando con la permutación determinada, por lo tanto, ∀ i . 0 ≤ i < rank ⇒ input_dimensions[permutation[i]] = output_dimensions[i].

Es lo mismo que Reshape(operando, permutación, Permute(permutación, operando.shape.dimensions)).

TriangularSolve

Consulta también XlaBuilder::TriangularSolve.

Resuelve sistemas de ecuaciones lineales con matrices de coeficiente triangular inferior o superior mediante sustitución hacia adelante o atrás. Cuando se transmite a través de dimensiones principales, esta rutina resuelve uno de los sistemas de matriz op(a) * x = b o x * op(a) = b para la variable x, dadas a y b, en las que op(a) es op(a) = a, op(a) = Transpose(a) o op(a) = Conj(Transpose(a)).

TriangularSolve(a, b, left_side, lower, unit_diagonal, transpose_a)

Argumentos Tipo Semántica
a XlaOp un array de rango > 2 de un tipo de punto flotante o complejo con la forma [..., M, M].
b XlaOp un array de rango > 2 del mismo tipo con la forma [..., M, K] si left_side es verdadero, [..., K, M] de lo contrario
left_side bool indica si se debe resolver un sistema con el formato op(a) * x = b (true) o x * op(a) = b (false).
lower bool si se debe usar el triángulo superior o inferior de a.
unit_diagonal bool si es true, se supone que los elementos diagonales de a son 1 y no se accede a ellos.
transpose_a Transpose si se debe usar a tal como está, se transpone o toma su transposición conjugada.

Los datos de entrada solo se leen desde el triángulo inferior/superior de a, según el valor de lower. Se ignoran los valores del otro triángulo. Los datos de salida se muestran en el mismo triángulo; los valores en el otro triángulo se definen en la implementación y pueden ser de cualquier tipo.

Si la clasificación de a y b es mayor que 2, se tratan como lotes de matrices y todas, excepto las 2 dimensiones menores, son dimensiones de lote. a y b deben tener las mismas dimensiones de lote.

Tupla

Consulta también XlaBuilder::Tuple.

Es una tupla que contiene un número variable de controladores de datos, cada uno de los cuales tiene su propia forma.

Esto es análogo a std::tuple en C++. Conceptualmente:

let v: f32[10] = f32[10]{0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
let s: s32 = 5;
let t: (f32[10], s32) = tuple(v, s);

Se puede deconstruir (se puede acceder a las tuplas) a través de la operación GetTupleElement.

Mientras que

Consulta también XlaBuilder::While.

While(condition, body, init)

Argumentos Tipo Semántica
condition XlaComputation XlaComputation de tipo T -> PRED, que define la condición de finalización del bucle.
body XlaComputation XlaComputation de tipo T -> T que define el cuerpo del bucle.
init T Valor inicial para el parámetro de condition y body.

Ejecuta secuencialmente body hasta que condition falla. Esto es similar a un bucle while típico en muchos otros idiomas, excepto por las diferencias y restricciones que se indican a continuación.

  • Un nodo While muestra un valor de tipo T, que es el resultado de la última ejecución de body.
  • La forma del tipo T se determina de forma estática y debe ser la misma en todas las iteraciones.

Los parámetros T de los cálculos se inicializan con el valor init en la primera iteración y se actualizan automáticamente al resultado nuevo desde body en cada iteración posterior.

Un caso de uso principal del nodo While es implementar la ejecución repetida del entrenamiento en redes neuronales. A continuación, se muestra un seudocódigo simplificado con un gráfico que representa el cálculo. Puedes encontrar el código en while_test.cc. El tipo T en este ejemplo es un Tuple que consta de un int32 para el recuento de iteraciones y un vector[10] para el acumulador. Para 1,000 iteraciones, el bucle agrega un vector constante al acumulador.

// Pseudocode for the computation.
init = {0, zero_vector[10]} // Tuple of int32 and float[10].
result = init;
while (result(0) < 1000) {
  iteration = result(0) + 1;
  new_vector = result(1) + constant_vector[10];
  result = {iteration, new_vector};
}