Semántica de operaciones

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

Nota sobre la nomenclatura: El tipo de datos generalizados con el que trata 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, se usa array 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 pueden vincularse entre operaciones de efecto secundario para aplicar el orden. AfterAll se puede usar como una unión de tokens para ordenar una operación después de un conjunto de operaciones.

AfterAll(operands)

Argumentos	Tipo	Semántica
operands	XlaOp	Cantidad de tokens variadic

AllGather

Consulta también XlaBuilder::AllGather.

Realiza la concatenación en todas las réplicas.

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

Argumentos	Tipo	Semántica
operand	XlaOp	Es el array que se concatenará en las 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 del 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 cuyo caso todas las réplicas pertenecen a un solo grupo, ordenado 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. Necesitamos esto en los casos en que replica_groups esté vacío.
• 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 salida es la forma de entrada con el all_gather_dim aumentado shard_count veces. Por ejemplo, si hay dos réplicas y el operando tiene los valores [1.0, 2.5] y [3.0, 5.25], respectivamente, en las dos réplicas, el valor de salida de esta operación 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 todas las réplicas.

AllReduce(operand, computation, replica_group_ids, channel_id)

Argumentos	Tipo	Semántica
operand	XlaOp	Es 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 arrays, la reducción total se realiza 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 la 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 los valores [1.0, 2.5] y [3.0, 5.25], respectivamente, en las dos réplicas, el valor de salida de esta operación y el cálculo de suma será [4.0, 7.75] en ambas réplicas. Si la entrada es una tupla, el resultado también es una tupla.

El cálculo del resultado de AllReduce requiere tener una entrada de cada réplica, por lo que, si una réplica ejecuta un nodo AllReduce más veces que otra, la primera réplica esperará para siempre. Dado que todas las réplicas ejecutan el mismo programa, no hay muchas formas de que eso suceda, pero es posible cuando la condición de un bucle while depende de los datos del feed y los datos que se insertan hacen que el bucle while 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. Tiene dos fases:

1. La fase de dispersión. En cada núcleo, el operando se divide en split_count cantidad de bloques a lo largo de split_dimensions, y los bloques se dispersan a 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 junto con el concat_dimension.

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

• replica_groups: Cada ReplicaGroup contiene una lista de 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á otra operación AllToAll dentro de las réplicas 4, 5 y 0, y el orden de concatenación también será 4, 5 y 0. Si replica_groups está vacío, todas las réplicas pertenecen a un grupo, en el orden de concatenación de su aparición.

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 n dimensiones
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 asigna un nombre a la dimensión a lo largo de la cual se concatenan los bloques divididos
split_count	int64	La cantidad de núcleos que participan en esta operación. Si replica_groups está vacío, esta debe ser la cantidad de réplicas; de lo contrario, debe ser igual a la cantidad de réplicas en cada grupo.
replica_groups	Vector ReplicaGroup	Cada grupo contiene una lista de IDs de réplicas.

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 la operación Alltoall. En cada núcleo, el operando se divide en 4 partes a lo largo de la dimensión 0, por lo que cada parte tiene la forma f32[4,4]. Las 4 partes se distribuyen en todos los núcleos. Luego, cada núcleo concatena las partes recibidas en la dimensión 1, en el orden del núcleo 0-4. Por lo tanto, el resultado en cada núcleo tiene la forma f32[16,4].

BatchNormGrad

Consulta también XlaBuilder::BatchNormGrad y el artículo original sobre la normalización por lotes para obtener una descripción detallada del algoritmo.

Calcula los gradientes de la normalización por lotes.

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

Argumentos	Tipo	Semántica
operand	XlaOp	Array de n dimensiones que se normalizará (x)
scale	XlaOp	Arreglo de 1 dimensión (\(\gamma\))
mean	XlaOp	Array de 1 dimensión (\(\mu\))
variance	XlaOp	Arreglo 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 componente de la dimensión de componentes (feature_index es el índice de la dimensión de componentes 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 con las siguientes fórmulas (suponiendo un array de 4 dimensiones como operand y con el índice de dimensión de 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 a la entrada operand ($\nabla x$)
grad_scale	XlaOp	gradiente con respecto a la entrada scale ($\nabla \gamma$)
grad_offset	XlaOp	gradiente con respecto a la entrada offset ($\nabla \beta$)

BatchNormInference

Consulta también XlaBuilder::BatchNormInference y el artículo original sobre la normalización por lotes para obtener una descripción detallada del algoritmo.

Normaliza un array en las dimensiones espaciales y de lote.

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

Argumentos	Tipo	Semántica
operand	XlaOp	matriz n dimensional para 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 atributo (feature_index es el índice de la dimensión del atributo en operand), la operación calcula la media y la varianza en 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 de componentes en operand.

BatchNormInference equivale a llamar a BatchNormTraining sin calcular mean y variance para cada lote. En su lugar, usa las entradas mean y variance como valores estimados. El propósito de esta operación es reducir la latencia en la inferencia; de ahí el nombre de BatchNormInference.

El resultado es un array normalizado de n dimensiones con la misma forma que el objeto 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 las dimensiones espaciales y de lote.

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

Argumentos	Tipo	Semántica
operand	XlaOp	Array de n dimensiones que se normalizará (x)
scale	XlaOp	Array de 1 dimensión (\(\gamma\))
offset	XlaOp	Arreglo 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 en la dimensión de atributos (feature_index es el índice de la dimensión de atributos en operand), la operación calcula la media y la varianza en 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 atributo en operand.

El algoritmo funciona de la siguiente manera para cada lote en operand \(x\) que contiene 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 de atributos: \(\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 de épsilon, que suele ser 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	Arreglo de 1 dimensión (\(\mu\))
batch_var	XlaOp	Array de 1 dimensión (\(\sigma^2\))

batch_mean y batch_var son momentos calculados en las dimensiones espaciales y del lote con las fórmulas anteriores.

BitcastConvertType

Consulta también XlaBuilder::BitcastConvertType.

Al igual que un tf.bitcast en TensorFlow, realiza una operación de transmisión de bits por elemento de una forma de datos a una forma objetivo. El tamaño de entrada y salida debe 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 convierte en cuatro elementos s8. La transmisión de bits se implementa como una transmisión de bajo nivel, por lo que las máquinas con diferentes representaciones de punto flotante darán resultados diferentes.

BitcastConvertType(operand, new_element_type)

Argumentos	Tipo	Semántica
operand	XlaOp	array de tipo T con dimensiones D
new_element_type	PrimitiveType	tipo U

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

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

Conversión de transmisión de bits a tipo primitivo de diferente ancho

La instrucción HLO BitcastConvert 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. Como toda la operación es conceptualmente una transmisió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 B > B', la forma de salida obtiene una nueva dimensión secundaria 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 elemento.

Transmisión

Consulta también XlaBuilder::Broadcast.

Agrega dimensiones a un array duplicando los datos en él.

Broadcast(operand, broadcast_sizes)

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

Las dimensiones nuevas se insertan a la izquierda, es decir, si broadcast_sizes tiene valores {a0, ..., aN} y la forma del operando tiene dimensiones {b0, ..., bM}, la forma del resultado tiene dimensiones {a0, ..., aN, b0, ..., bM}.

Las nuevas dimensiones indexan copias del operando, es decir,

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 del resultado serán 2.0f.

BroadcastInDim

Consulta también XlaBuilder::BroadcastInDim.

Expande el tamaño y el rango de un array duplicando 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 objetivo
broadcast_dimensions	ArraySlice<int64>	A qué dimensión de la forma objetivo corresponde cada dimensión de la forma del operando

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

El 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 el tamaño 1 o tener el mismo tamaño que la dimensión en la forma de salida a la que se asignan. Las dimensiones restantes se completan con dimensiones de tamaño 1. Luego, la transmisión de dimensiones degeneradas se transmite a lo largo de estas dimensiones para llegar a la forma de salida. La semántica se describe en detalle en la página de transmisión.

Llamar

Consulta también XlaBuilder::Call.

Invoca un cálculo con los argumentos proporcionados.

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 aridad y los tipos de args deben coincidir con los parámetros de computation. No puede tener args.

Cholesky

Consulta también XlaBuilder::Cholesky.

Calcula la decomposition de Cholesky de un lote de matrices simétricas (hermíticas) positivas y definidas.

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 del triángulo inferior o 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 del otro triángulo se definen en la implementación y pueden ser cualquier cosa.

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 por lotes.

Si a no es un valor simétrico (hermitiano) definido de forma positiva, el resultado está definido por la implementación.

Abrazadera

Consulta también XlaBuilder::Clamp.

Limita un operando dentro del rango entre un valor mínimo y 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

Dado un operando y valores mínimos y máximos, muestra el operando si está dentro del 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 el operando está por encima de este rango. Es decir, clamp(a, x, b) = min(max(a, x), b).

Los tres arrays deben tener la misma forma. Como alternativa, como forma restringida de emisión, min o max pueden ser un escalar de tipo T.

Ejemplo con min escalares y max:

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 dimensión.

Collapse(operand, dimensions)

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

La operación de colapso reemplaza el subconjunto determinado de las dimensiones del operando por una sola dimensión. Los argumentos de entrada son un array arbitrario de tipo T y un vector de índices de dimensión constante en el tiempo de compilación. Los índices de dimensión deben ser un subconjunto consecutivo en orden (números de dimensión de baja a alta) de las dimensiones de T. Por lo tanto, {0, 1, 2}, {0, 1} o {1, 2} son 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 que las que reemplazan, con un 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 de variación más lenta (más principal) en el nido de bucles que contrae estas dimensiones, y el número de dimensión más alto es la variación más rápida (más menor). Consulta el operador tf.reshape si necesitas un orden de colapso 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 réplicas de datos entre réplicas.

CollectivePermute(operand, source_target_pairs)

Argumentos	Tipo	Semántica
operand	XlaOp	Array de entrada de n dimensiones
source_target_pairs	Vector <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 source_target_pair:

• Ningún par debe 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, el resultado en esa réplica es un tensor que consta de 0 con la misma forma que la entrada.

Concatenate

Consulta también XlaBuilder::ConcatInDim.

Concatenar compone un array a partir de varios operandos de array. El array tiene el mismo rango que cada uno de los operandos del array de entrada (que deben tener el mismo rango entre sí) y contiene los argumentos en el orden en que se especificaron.

Concatenate(operands..., dimension)

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

Con la excepción de dimension, todas las dimensiones deben ser iguales. Esto se debe a que XLA no admite arrays “irigentes”. También 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 de 1 dimensión:

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

Ejemplo de 2 dimensiones:

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:

Condicional

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 recibir un solo argumento de tipo \(T_0\) y se invocará con true_operand, que debe ser del mismo tipo. false_computation debe aceptar 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á una de 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 un S32 que es < 0 o >= N, entonces branch_computations[N-1] se ejecuta como la rama predeterminada.

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

Ten en cuenta que solo se ejecutará uno de los 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. El padding SAME agrega ceros a la entrada (lhs) para que el resultado tenga la misma forma que la entrada cuando no se tenga en cuenta el paso. El padding VALID simplemente significa que no hay padding.

ConvWithGeneralPadding (convolución)

Consulta también XlaBuilder::ConvWithGeneralPadding.

Calcula una convolución del tipo que se usa en las redes neuronales. Aquí, se puede considerar una convolución como una ventana de n dimensiones que se mueve a través de un área base de n dimensiones y se realiza un cálculo para cada posición posible de la ventana.

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

Supongamos que n es la cantidad de dimensiones espaciales. El argumento lhs es un array de rango n+2 que describe el área base. Esto se denomina 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 son, 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 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, el kernel o la ventana convolucional. Las dimensiones son las siguientes:

• output-z: Es la dimensión z del resultado.
• input-z: El tamaño de esta dimensión por feature_group_count debe ser 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 paso de la ventana de convolución en las dimensiones espaciales. Por ejemplo, si el paso en la primera dimensión espacial es 3, la ventana solo se puede colocar en coordenadas donde el primer índice espacial sea divisible por 3.

El argumento padding especifica la cantidad de relleno de cero que se aplicará al área de base. La cantidad de relleno puede ser negativa. El valor absoluto del relleno negativo indica la cantidad de elementos que se deben quitar de la dimensión especificada antes de realizar la convolución. padding[0] especifica el padding para la dimensión y y padding[1] especifica el padding para la dimensión x. Cada par tiene el padding bajo como primer elemento y el padding alto como segundo elemento. 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 de 3 ceros a la derecha en la segunda dimensión espacial. El uso del padding equivale a insertar esos mismos valores 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 la izq. y la der., respectivamente, en cada dimensión espacial. Si el factor de dilatación en una dimensión espacial es d, se colocan implícitamente d-1 orificios entre cada una de las entradas de esa dimensión, lo que aumenta el tamaño del array. Los agujeros se completan con un valor no operativo, que para la convolución significa cero.

La dilatación del lado derecho también se denomina convolución atrous. Para obtener más información, consulta tf.nn.atrous_conv2d. La dilatación de la lhs también se conoce como convolución transposta. Para obtener más información, consulta tf.nn.conv2d_transpose.

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

Para la convolución en profundidad, el argumento feature_group_count se establecería en la dimensión de atributos de entrada, y el filtro se modificarí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 (de entrada). Si batch_group_count es mayor que 1, significa que la dimensión del lote de salida debe tener un tamaño de input batch / batch_group_count. El batch_group_count debe ser un divisor del tamaño de la función de salida.

La forma de salida tiene estas 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 ubicación válida de la ventana convolucional.

En la figura anterior, se muestra cómo funciona el campo batch_group_count. De hecho, divide cada lote de lhs en grupos batch_group_count y hacemos lo mismo con los atributos de salida. Luego, para cada uno de estos grupos, hacemos 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 (espaciales y de componentes) sigue siendo la misma.

Las posiciones válidas de la ventana de convolución se determinan según los pasos y el tamaño del área base después del padding.

Para describir lo que hace una convolución, considera una convolución en 2D y elige algunas coordenadas batch, z, y, x fijas en el resultado. Luego, (y,x) es la 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 del área de base, en la que cada punto de 2D está asociado a un vector de 1D, por lo que obtenemos un cuadro 3D. Desde el kernel de convolución, como corregimos la coordenadas de salida z, también tenemos un cuadro 3D. Las dos cajas 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 de convolución se usa; hay un cuadro 3D separado de valores que se usa para cada coordenada output-z. Por lo tanto, puedes considerarlo como 5 convoluciónes separadas con un filtro diferente para cada una de ellas.

Este es el seudocódigo de una convolución 2D con padding y stride:

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 un static_cast por elemento en C++, realiza una operación de conversión por elemento de una forma de datos a una forma de destino. Las dimensiones deben coincidir, y la conversión es por elemento; p. ej., los elementos s32 se convierten en elementos f32 a través de una rutina de conversión de s32 a f32.

ConvertElementType(operand, new_element_type)

Argumentos	Tipo	Semántica
operand	XlaOp	array de tipo T con dimensiones D
new_element_type	PrimitiveType	tipo U

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

Una conversión como T=s32 a U=f32 realizará una rutina normalización de conversión de número entero a flotante, por ejemplo, de redondo 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 procesamiento de suma.

CustomCall

Consulta también XlaBuilder::CustomCall.

Llama a una función proporcionada por el usuario dentro de un procesamiento.

CustomCall(target_name, args..., shape)

Argumentos	Tipo	Semántica
target_name	string	Es el nombre de la función. Se emitirá una instrucción de llamada que se orientará 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 de la salida 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])

A continuación, se muestra 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	Salida	Semántica
vector [n] dot vector [n]	escalar	producto punto de vectores
matriz [m × k] vector dot [k]	vector [m]	multiplicación matriz-vector
matriz [m x k] dot matriz [k x n]	matriz [m × n]	multiplicación de matrices

La operación realiza la suma de productos en la segunda dimensión de lhs (o la primera si tiene rango 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 punto entre vectores, multiplicaciones de vectores/matrices o multiplicaciones de matrices/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 contratación y por lotes

Es similar a Dot, pero permite especificar números de dimensión de lotes y contratos para lhs y rhs.

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

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

No es necesario que los números de dimensiones de contratación asociados de lhs y rhs sean los mismos, 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 por lotes 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 2 × 2):

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	Salida	Semántica
[b0, m, k] dot [b0, k, n]	[b0, m, n]	matmul por lotes
[b0, b1, m, k] dot [b0, b1, k, n]	[b0, b1, m, n]	matmul por lotes

Por lo tanto, el número de dimensión resultante comienza con la dimensión de lote, luego con la dimensión lhs sin contrato o sin lote y, por último, con la dimensión rhs sin contrato o sin lote.

DynamicSlice

Consulta también XlaBuilder::DynamicSlice.

DynamicSlice extrae un subconjunto del array 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 final de los intervalos de porciones exclusivas en cada dimensión: [inicio, inicio + tamaño]. La forma de start_indices debe ser de rango == 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 dimensional N de tipo T
start_indices	secuencia de N XlaOp	Es una 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>	Es una lista de N números enteros que contienen el tamaño de la porción para cada dimensión. Cada valor debe ser estrictamente mayor que cero, y el inicio + tamaño debe ser menor o igual que el tamaño de la dimensión para evitar el ajuste del tamaño de la dimensión modulo.

Los índices de fragmento efectivos se calculan aplicando la siguiente transformación para cada índice i en [1, N) antes de realizar la divisió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 array del operando. Si el fragmento está dentro de los límites antes de que se aplique la transformación, esta no tendrá efecto.

Ejemplo de 1 dimensión:

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 de 2 dimensiones:

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 array de entrada operand, con una porción update reemplazada en start_indices. La forma de update determina la forma del subarray del resultado que se actualiza. La forma de start_indices debe ser de clasificación == 1, con un tamaño de dimensión igual a la clasificación de operand.

DynamicUpdateSlice(operand, update, start_indices)

Argumentos	Tipo	Semántica
operand	XlaOp	Array de N dimensiones de tipo T
update	XlaOp	Arreglo N dimensional de tipo T que contiene la actualización de los segmentos. Cada dimensión de la forma de actualización debe ser estrictamente mayor que cero, y el inicio + la actualización debe ser menor o igual que el tamaño del operando para cada dimensión para evitar generar índices de actualización fuera de rango.
start_indices	secuencia de N XlaOp	Es una 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 fragmento efectivos se calculan aplicando la siguiente transformación para cada índice i en [1, N) antes de realizar la divisió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 con respecto al array del operando. Si el fragmento está dentro de los límites antes de que se aplique la transformación, esta no tendrá efecto.

Ejemplo de 1 dimensión:

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 de 2 dimensiones:

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 por elemento.

Op(lhs, rhs)

Donde Op es uno de los siguientes operadores: Add (suma), Sub (resta), Mul (multiplicación), Div (división), Pow (potencia), Rem (residuo), Max (máximo), Min (mínimo), And (AND lógico), Or (O lógico), Xor (XOR lógico), ShiftLeft (Mayúsculas a la izquierda), ShiftRightArithmetic (Mayúsculas a la derecha aritméticas), ShiftRightLogical (Mayúsculas a la derecha lógicas), Atan2 (arcotangente de 2 argumentos) o Complex (combina partes reales e imaginarias en un número complejo).

Argumentos	Tipo	Semántica
lhs	XlaOp	Operando del lado izquierdo: Array de tipo T
rhs	XlaOp	Operando del lado derecho: Array de tipo T

Las formas de los argumentos deben ser similares o compatibles. Consulta la documentación sobre transmisiones para saber qué significa que las formas sean compatibles. El resultado de una operación tiene una forma que es el resultado de la transmisión de los dos arrays de entrada. En esta variante, no se admiten operaciones entre arrays de diferentes rangos, 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 siempre es menor que el valor absoluto del divisor.

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

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

Op(lhs, rhs, broadcast_dimensions)

En el que Op es igual que el anterior. Esta variante de la operación debería usarse 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 la clasificación del operando de rango inferior hasta la clasificación 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 no asignadas de la forma expandida se completan con dimensiones de tamaño uno. La transmisión de degeneración de dimensión luego transmite las formas a lo largo de estas dimensiones degeneradas para igualar las formas de ambos operandos. La semántica se describe en detalle en la página de transmisión.

Operaciones de comparación en elementos

Consulta también XlaBuilder::Eq.

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

Op(lhs, rhs)

En el ejemplo anterior, Op es Eq (igual que), Ne (no es igual que), Ge (mayor o igual que), Gt (mayor que), Le (menor o igual que) y Lt (menor que). Otro conjunto de operadores, EqTotalOrder, NeTotalOrder, GeTotalOrder, GtTotalOrder, LeTotalOrder y LtTotalOrder, proporcionan las mismas funcionalidades, excepto que, además, admiten un orden total sobre los números de punto flotante, ya que aplican -NaN < -Inf < -Finite < -0 < +0 < +Finite < +Inf < +NaN.

Argumentos	Tipo	Semántica
lhs	XlaOp	Operando del lado izquierdo: Array de tipo T
rhs	XlaOp	Operando del lado derecho: Array de tipo T

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

Existe una variante alternativa con compatibilidad con transmisiones de diferentes rangos para estas operaciones:

Op(lhs, rhs, broadcast_dimensions)

En el que Op es igual que el anterior. Esta variante de la operación se debe usar para operaciones de comparación entre arrays de diferentes rangos (como agregar una matriz a un vector).

El operando broadcast_dimensions adicional es un fragmento de números enteros que especifica las dimensiones que se usarán para transmitir los operandos. La semántica se describe en detalle en la página de transmisión.

Funciones unarias en términos de elementos

XlaBuilder admite las siguientes funciones unarias por elemento:

Abs(operand) Función abs por elemento x -> |x|.

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

Ceil(operand) Techo por elemento x -> ⌈x⌉.

Clz(operand) Cuenta los ceros iniciales por elemento.

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

Erf(operand) Función de error por elemento x -> erf(x), donde

\(\text{erf}(x) = \frac{2}{\sqrt{\pi} }\int_0^x e^{-t^2} \, dt\).

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

Expm1(operand) Exponencial natural por elemento menos uno x -> e^x - 1.

Floor(operand) Precio mínimo por elemento x -> ⌊x⌋.

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

IsFinite(operand) Comprueba si cada elemento de operand es finito, es decir, no es infinito positivo o 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 del elemento x -> ln(x).

Log1p(operand) Logaritmo natural desplazado por elemento x -> ln(1+x).

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

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

Not(operand) Lógico de elemento inteligente no x -> !(x).

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

Real(operand) Es la parte real en relación con los elementos de una forma compleja (o real). x -> real(x). Si el operando es de tipo de punto flotante, muestra el mismo valor.

Round(operand) Redondeo por elemento, se aleja de cero.

RoundNearestEven(operand) Redondeo por elemento, se vincula al número par más cercano.

Rsqrt(operand) El recíproco en elementos de la operación raíz cuadrada x -> 1.0 / sqrt(x).

Sign(operand) Operación de signo en elementos x -> sgn(x) en la que

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

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

Sin(operand) Seno por elemento x -> sin(x).

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

Tan(operand) Tangente por elemento x -> tan(x).

Tanh(operand) Tangente hiperbólica en elemento x -> tanh(x).

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 genera 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 directas 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	El array que transformamos de Fourier.
fft_type	FftType	Consulta la tabla que se encuentra a continuación
fft_length	ArraySlice<int64>	Las longitudes del dominio de tiempo de los ejes que se transforman. Esto es necesario, en particular, para que IRFFT ajuste el tamaño del 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	FFT de complejo a complejo de reenvío. La forma no cambia.
IFFT	FFT inverso de complejo a complejo. La forma no cambia.
RFFT	Reenvía la FFT de real a complejo. 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 y 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 del conjugado inverso de las entradas de 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, para los casos de real a complejo y complejo a real, la transformación del eje más interno se realiza (de manera efectiva) primero (RFFT; la última para IRFFT), por lo que el eje más interno es el que cambia de tamaño. Las demás transformaciones de ejes serán de complejo a complejo.

Detalles de implementación

Los FFT de la CPU están respaldados por el TensorFFT de Eigen. La FFT de GPU 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 que estamos recopilando datos.
start_indices	XlaOp	Es un array 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 de inicio. Consulta la siguiente información para obtener una descripción detallada.
offset_dims	ArraySlice<int64>	Es el conjunto de dimensiones en la forma de salida que se compensa en un array cortado del operando.
slice_sizes	ArraySlice<int64>	slice_sizes[i] son los límites para 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 tener un tamaño de 1.
start_index_map	ArraySlice<int64>	Mapa en el que se describe cómo asignar índices de start_indices a índices legales en operandos.
indices_are_sorted	bool	Indica si se garantiza que el llamador ordene los índices.

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

El resultado es un array de rango batch_dims.size + offset_dims.size.

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 tiene una dimensión 1 final (es decir, si start_indices tenía la forma [6,7] y index_vector_dim es 2, entonces consideramos implícitamente que la forma de start_indices es [6,7,1]).

Los límites del array de salida a lo largo de 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 algún k), elegimos los límites de dimensión correspondientes de start_indices.shape, omitiendo index_vector_dim (es decir, elegimos start_indices.shape.dims[k] si k < index_vector_dim y start_indices.shape.dims[k+1] de lo contrario).

2. Si i está presente en offset_dims (es decir, es igual a offset_dims[k] para algún k), elegimos el límite correspondiente de slice_sizes después de tener en cuenta collapsed_slice_dims (es decir, elegimos adjusted_slice_sizes[k] donde adjusted_slice_sizes es slice_sizes con los límites en los índices collapsed_slice_dims quitados).

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

1. Sea G = { Out[k] para k en batch_dims }. Usa G para cortar un vector S de modo que S[i] = start_indices[Combine(G, i)] donde 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, S = start_indices.

2. Crea un índice inicial, Sin, en operand con S. Para ello, dispérsalo con start_index_map. Más precisamente:

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

   2. Sin[_] = 0 en caso contrario.

3. Crea un índice Oin en operand dispersando los índices en las dimensiones de offset en Out según 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 en caso contrario.

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

remapped_offset_dims es una función monótona con dominio [0, offset_dims.size) y rango [0, operand.rank) \ collapsed_slice_dims. Por lo tanto, si, p. ej., Si offset_dims.size es 4, operand.rank es 6 y collapsed_slice_dims es {0, 2}, entonces remapped_offset_dims es {0→1, 1→3, 2→4, 3→5}.

Si indices_are_sorted se establece como verdadero, XLA puede suponer que el usuario ordena los start_indices (en orden ascendente, después de dispersar sus valores según start_index_map). Si no es así, la semántica se define en la implementación.

Descripción informal y ejemplos

De manera informal, cada índice Out en el array de salida corresponde a un elemento E en el array de operandos, que se calcula de la siguiente manera:

• Usamos las dimensiones del lote en Out para buscar un índice inicial desde 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.

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

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

• Usamos las dimensiones de compensación 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 establece en start_indices.rank - 1 en todos los ejemplos que se muestran a continuación. Los valores más interesantes para index_vector_dim no cambian la operación de forma fundamental, pero hacen que la representación visual sea más engorrosa.

Para obtener una idea de cómo se unen todo lo anterior, veamos un ejemplo que reúne 5 rebanadas de la forma [8,6] de un array [16,11]. La posición de una porción en el array [16,11] se puede representar como un vector de índices de 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 índice 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 índices de recopilación con G. El elemento del array de salida en el índice [G,O0,O1] es el elemento del array de entrada en el índice [X+O0,Y+O1].

slice_sizes es [8,6], que decide el rango de O₀ y O₁, 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 usa un array de "índices de recopilación" de forma [4,5,2] traduciría los índices de la siguiente manera:

Una vez más, 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 antes de las siguientes maneras:

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

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 fragmento de 1. Dado que tienen un tamaño de fragmento de 1, el único índice válido para ellos es 0, y su eliminación no genera ambigüedad.

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

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

G0 y G1 se usan para cortar un índice de inicio del array de índices de recopilación como de costumbre, excepto que el índice de inicio solo tiene un 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 "Gather Index Mapping" (start_index_map en la descripción formal) y "Offset Mapping" (remapped_offset_dims en la descripción formal) en [X,0] y [0,O0], respectivamente, que suman [X,O0]. En otras palabras, el índice de salida [G0,G1,O0] se asigna al índice de entrada [GatherIndices[G0,G1,0],O0], que nos brinda la semántica de tf.gather_nd.

slice_sizes para este caso es [1,11]. De manera intuitiva, esto significa que cada índice X en el array de índices de recopilación 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 n dimensiones
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	Es un array de entrada de n dimensiones.
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.

Como resultado, pasa el operando, y el compilador realiza un seguimiento de la dimensión dinámica.

Las operaciones de reducción descendente ignorarán los valores con padding.

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.

Indexa 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.

In-feed

Consulta también XlaBuilder::Infeed.

Infeed(shape)

Argumento	Tipo	Semántica
shape	Shape	Es la forma de los datos que se leen de la interfaz de In-feed. El campo de diseño de la forma debe establecerse para que coincida con el diseño de los datos que se envían al dispositivo. De lo contrario, su comportamiento no está definido.

Lee un solo elemento de datos de la interfaz de transmisión de in-feed implícita del dispositivo, interpreta los datos como la forma y el diseño determinados, y muestra un XlaOp de los datos. Se permiten varias operaciones de In-feed en un cálculo, pero debe haber un orden total entre las operaciones de In-feed. 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)
}

No se admiten las formas de tupla anidadas. Para una forma de tupla vacía, la operación de Infeed es, en realidad, una operación sin efecto y continúa sin leer ningún dato del Infeed 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 que tiene una forma especificada y contiene valores que comienzan en cero y aumentan en uno a lo largo de la dimensión especificada. Para los tipos de números 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 número de punto flotante.

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

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 ]]

Mapa

Consulta también XlaBuilder::Map.

Map(operands..., computation)

Argumentos	Tipo	Semántica
operands	secuencia de N XlaOp	N arreglos 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	int64 array	array de dimensiones del mapa

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

La función asignada es un cálculo arbitrario con la restricción de que tiene N entradas de tipo escalar T y una sola salida con el tipo S. El resultado 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) en cada índice (multidimensional) de los آرا￳os de entrada para producir el array de salida.

OptimizationBarrier

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

Garantiza que todas las entradas se evalúen antes que cualquier operador que dependa de los resultados 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	Es el valor del padding en ambos bordes (bajo, alto) y entre los elementos de cada dimensión.

Expande el array operand determinado agregando padding alrededor del array y entre los elementos del array con el padding_value determinado. padding_config especifica la cantidad de padding de borde y el padding interior para 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 superior (junto al índice más alto) de cada dimensión, respectivamente. La cantidad de padding de 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 en cada dimensión; no puede ser negativo. El padding interno se produce de manera lógica antes del padding del borde, por lo que, en el caso del padding en el borde negativo, los elementos se quitan del operando con padding interior.

Esta operación no realiza ninguna acción si los pares de padding de borde son todos (0, 0) y los valores de padding interior son todos 0. En la siguiente imagen, se muestran ejemplos de diferentes valores de edge_padding y interior_padding para un array de dos dimensiones.

Recv

Consulta también XlaBuilder::Recv.

Recv(shape, channel_handle)

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

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

La API cliente de la operación Recv representa la comunicación síncrona. Sin embargo, la instrucción se divide internamente 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 para esperar que se complete la transferencia de datos. El contexto es una tupla de {receive buffer (shape), request identifier (U32)} y solo puede ser utilizada por una instrucción RecvDone.

RecvDone(HloInstruction context)

Dado 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 arreglos de tipos T_0, ..., T_{N-1}
init_values	Secuencia de N XlaOp	N escalares de tipos T_0, ..., T_{N-1}.
computation	XlaComputation	procesamiento de tipo T_0, ..., T_{N-1}, T_0, ..., T_{N-1} -> Collate(T_0, ..., T_{N-1}).
dimensions	int64 array	Es un array desordenado de dimensiones que se reducirán.

Aquí:

• N debe ser mayor o igual que 1.
• El procesamiento debe ser "aproximadamente" asociativo (consulta a continuación).
• Todos los arrays de entrada deben tener las mismas dimensiones.
• Todos los valores iniciales deben formar una identidad en 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. El rango de cada array 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.

Los diferentes backends pueden volver a asociar el procesamiento de reducción. Esto puede generar diferencias numéricas, ya que algunas funciones de reducción, como la adición, no son asociativas para los números de punto flotante. Sin embargo, si el rango de los datos es limitado, la adición de punto flotante es lo suficientemente cercana a ser asociativa para la mayoría de los usos prácticos.

Ejemplos

Cuando se realiza una reducción de una dimensión en un solo array 1D con los valores [10, 11, 12, 13], con la función de reducción f (esto es computation), el cálculo podría calcularse de la siguiente manera:

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

pero también hay 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 pseudocódigo de cómo se podría implementar la reducción, usando la suma como el cálculo de 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 la reducción de dimensiones 0 o 1 con una función "agregar":

Ten en cuenta que ambos resultados de reducción son arrays de 1D. El diagrama muestra uno como columna y otro como fila solo por conveniencia visual.

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

De manera similar al ejemplo de 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 combinaron 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 combinaron en un escalar:

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

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

También podemos reducir varias dimensiones. Las dimensiones de reducción de suma 0 y 1 producen el array [20, 28, 36] de 1D.

La reducción del array 3D en todas sus dimensiones produce el escalar 84.

Reducción variadica

Cuando N > 1, la aplicación de la función reduce es un poco más compleja, ya que se aplica de forma simultánea a todas las entradas. Los operandos se proporcionan a la operación en el siguiente orden:

• Ejecución de un valor reducido para el primer operando
• …
• Ejecución de un valor reducido para el operando N
• Valor de entrada para el primer operando
• …
• Valor de entrada para el operando N

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 1 dimensión 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 arreglos de entrada V = Float[N], K = Int[N] de 1 dimensión y los valores iniciales I_V = Float, I_K = Int, el resultado f_(N-1) de reducir 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á en los arrays y se mostrará una tupla que contiene 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 volver al formato original. La cantidad de bits exponenciales y de mantisa en el formato de menor precisión puede especificarse de forma arbitraria, aunque es posible que todos los tamaños de bits no 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 de exponente en formato de menor precisión
mantissa_bits	int32	Cantidad de bits de mantisa en formato de precisión inferior

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

El formato de menor precisión debe tener al menos un bit exponente (para distinguir un valor cero de un infinito, ya que ambos tienen una mantisa cero) y debe tener un número no negativo de bits de mantisa. La cantidad de bits exponentes o mantissa puede superar 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 una operación AllReduce de manera eficaz y, luego, dispersa el resultado dividiéndolo en bloques shard_count a lo largo de scatter_dimension y la réplica i en el 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 para reducir en todas las réplicas.
computation	XlaComputation	Cálculo de reducción
scatter_dimension	int64	Es la dimensión que se dispersa.
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 del canal opcional para la comunicación entre módulos

• Cuando operand es una tupla de arrays, la reducción-dispersión se realiza en cada elemento de la tupla.
• replica_groups es una lista de los grupos de réplicas entre los que se realiza la reducción (el ID 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 Allreduce. 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. Cuando hay más de un grupo de réplicas, todos deben ser del mismo tamaño. Por ejemplo, replica_groups = {0, 2}, {1, 3} realiza una reducción entre las réplicas entre las réplicas 0 y 2, y entre 1 y 3, y, luego, dispersa el resultado.
• shard_count es el tamaño de cada grupo de réplicas. La necesitamos para los casos en los que replica_groups está vacío. 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 forma de entrada con el scatter_dimension reducido shard_count veces. Por ejemplo, si hay dos réplicas y el operando tiene los valores [1.0, 2.25] y [3.0, 5.25], respectivamente, en las dos 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 sola o una tupla de N arrays multidimensionales como resultado. Cada array de salida tiene la misma cantidad de elementos que la cantidad de posiciones válidas de la ventana. Una capa de agregación se puede expresar como una ReduceWindow. Al igual que con Reduce, el computation aplicado siempre 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 multidimensionaes 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	N valores de inicio para la reducción, uno para cada uno de los N operandos. Consulta Reduce para obtener más detalles.
computation	XlaComputation	Es una función de reducción de tipo T_0, ..., T_{N-1}, T_0, ..., T_{N-1} -> Collate(T_0, ..., T_{N-1}) que se aplica a los elementos de cada ventana de todos los operandos de entrada.
window_dimensions	ArraySlice<int64>	Es un array de números enteros para los valores de la dimensión de ventana.
window_strides	ArraySlice<int64>	array de números enteros para valores de zancada de la ventana
base_dilations	ArraySlice<int64>	array de números enteros para valores de dilatación de bases
window_dilations	ArraySlice<int64>	Es un array de números enteros para los valores de dilatación de la ventana.
padding	Padding	tipo de padding para la ventana (Padding::kSame, que agrega padding para tener la misma forma de salida que la entrada si el paso es 1, o Padding::kValid, que no usa padding y "detiene" la ventana una vez que ya no se ajusta)

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}).

En el siguiente código y la siguiente imagen, se muestra un ejemplo del uso de ReduceWindow. La entrada es una matriz de tamaño [4 × 6] y window_dimensions y window_stride_dimensions son [2 × 3].

// 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);

Un intervalo de 1 en una dimensión especifica que la posición de una ventana en la dimensión está a 1 elemento de su ventana adyacente. Para especificar que ninguna ventana se superponga entre sí, window_stride_dimensions debe ser igual a window_dimensions. En la siguiente figura, se ilustra el uso de dos valores de stride diferentes. El padding se aplica a cada dimensión de la entrada, y los cálculos son los mismos que si la entrada llegara con las dimensiones que tiene después del padding.

Para un ejemplo de padding no trivial, considera calcular el mínimo de la ventana reducida (el valor inicial es MAX_FLOAT) con la dimensión 3 y el paso 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]. El padding kSame primero agrega padding al array para que la forma después de la ventana de reducción sea la misma que la entrada para el paso uno agregando elementos iniciales en ambos lados, lo que genera [MAX_VALUE, 10000, 1000, 100, 10, 1, MAX_VALUE]. Ejecutar reduce-window en el array con padding opera en tres ventanas [MAX_VALUE, 10000, 1000], [1000, 100, 10], [10, 1, MAX_VALUE] y genera [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 análisis 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 signo en el intervalo [0, N), donde N es la cantidad de réplicas. Como todas las réplicas ejecutan el mismo programa, una llamada a ReplicaId() en el programa mostrará un valor diferente en cada réplica.

Reshape

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

Redefine las dimensiones de un array en 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 int64	orden en que se contraen las dimensiones
new_sizes	Vector int64	vector de tamaños de dimensiones nuevas

Conceptualmente, la reforma primero aplana un array y lo convierte en un vector unidimensional de valores de datos y, luego, define mejor este vector para obtener una nueva forma. Los argumentos de entrada son un array arbitrario de tipo T, un vector constante de tiempo de compilación de índices de dimensión y un vector constante de tiempo de compilación de tamaños de dimensión para el resultado. Si se proporcionan, los valores del vector dimension 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 es de la dimensión de variación más lenta (más principal) a la dimensión de variación más rápida (más secundaria) en el nido 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 en el índice 0 en new_sizes es el tamaño de la dimensión 0, el valor en el í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 las dimensiones del operando. Cuando se define el array colapsado en el array multidimensiona que define new_sizes, las dimensiones en new_sizes se ordenan de la variación más lenta (más importante) a la más rápida (menos importante).

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, reshape 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 (invertir)

Consulta también XlaBuilder::Rev.

Rev(operand, dimensions)

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

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

Un uso de la operación Rev es invertir el array de pesos de convolución a lo largo de las dos dimensiones de la ventana durante el cálculo del gradiente en las redes neuronales.

RngNormal

Consulta también XlaBuilder::RngNormal.

Construye un resultado de una forma determinada con números aleatorios generados siguiendo la distribución normal de \(N(\mu, \sigma)\) . Los parámetros \(\mu\) y \(\sigma\), y la forma de salida deben tener un tipo elemental de punto flotante. Además, los parámetros deben tener un valor escalar.

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	Es un escalar de tipo T que especifica la desviación estándar de los datos generados.
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 siguiendo la distribución uniforme en el intervalo \([a,b)\). Los parámetros y el tipo de elemento de salida deben ser un tipo booleano, un tipo de número entero o un tipo de punto flotante, y los tipos deben ser coherentes. Actualmente, 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	Un escalar de tipo T que especifica el límite inferior del intervalo
b	XlaOp	Un 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 llena de bits aleatorios uniformes con el algoritmo especificado (o el 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 de números aleatorios actual. Esta, la forma requerida y los valores válidos dependen del algoritmo que se use.

Se garantiza que el resultado sea una función determinística del estado inicial, pero no se garantiza que sea determinística entre los backends y las diferentes versiones del compilador.

RngBitGenerator(algorithm, key, shape)

Argumentos	Tipo	Semántica
algorithm	RandomAlgorithm	Es el algoritmo de PRNG que se usará.
initial_state	XlaOp	Estado inicial del algoritmo de PRNG.
shape	Shape	Es la forma de salida de los datos generados.

Valores disponibles para algorithm:

• rng_default: Algoritmo específico de backend con requisitos de forma específicos de backend.

• rng_three_fry: Algoritmo de PRNG basado en contadores ThreeFry. La forma initial_state es u64[2] con valores arbitrarios. Salmon et al. SC 2011. Números aleatorios paralelos: es tan fácil como 1, 2, 3.

• rng_philox: Algoritmo Philox para generar números aleatorios en paralelo. La forma initial_state es u64[3] con valores arbitrarios. Salmon et al. SC 2011. Números aleatorios paralelos: es tan fácil como 1, 2, 3.

Dispersión

La operación de dispersión de XLA genera una secuencia de resultados que son los valores del array de entrada operands, con varias porciones (en los índices especificados por scatter_indices) actualizadas con la secuencia de valores en updates usando 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 arreglos de tipos T_0, ..., T_N en los que se dispersarán.
scatter_indices	XlaOp	Es un array que contiene los índices iniciales de las porciones a las que se deben distribuir.
updates	Secuencia de N XlaOp	N arreglos de tipos T_0, ..., T_N updates[i] contiene los valores que se deben usar para la dispersión operands[i].
update_computation	XlaComputation	Cálculo que se usará 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 de scatter_indices que contiene los índices iniciales.
update_window_dims	ArraySlice<int64>	Es el conjunto de dimensiones en forma de updates que son dimensiones de ventana.
inserted_window_dims	ArraySlice<int64>	Es el conjunto de dimensiones de la ventana que se debe insertar en la forma updates.
scatter_dims_to_operand_dims	ArraySlice<int64>	Un mapa de 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]. Debe ser uno a uno y total.
indices_are_sorted	bool	Indica si se garantiza que el llamador ordene los índices.
unique_indices	bool	Indica si el llamador garantiza que los índices sean únicos.

Aquí:

• Se requiere que N sea mayor o igual que 1.
• operands[0], …, operands[N-1] deben tener las mismas dimensiones.
• updates[0], …, 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, consideramos implícitamente que scatter_indices tenga una dimensión 1 final.

Definimos update_scatter_dims de tipo ArraySlice<int64> como el conjunto de dimensiones en forma de 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 algunos k), el límite de la dimensión i en updates no debe exceder el límite correspondiente de operand después de contabilizar el inserted_window_dims (es decir, adjusted_window_bounds[k], en el que 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, es igual a update_scatter_dims[k] para algún k), el límite de la dimensión i en updates debe ser igual al límite correspondiente de scatter_indices, omitiendo index_vector_dim (es decir, scatter_indices.shape.dims[k], si k < index_vector_dim y scatter_indices.shape.dims[k+1] de lo contrario).

• update_window_dims debe estar en orden ascendente, no tener números de dimensión que se repitan y estar en el rango [0, updates.rank).

• inserted_window_dims debe estar en orden ascendente, no debe tener números de dimensión repetidos y debe 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 U determinado en cada array updates, el índice I correspondiente en el array operands correspondiente en el que se debe aplicar esta actualización se calcula de la siguiente manera:

1. Sea 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) inserta b en las posiciones index_vector_dim en A.
2. Crea un índice Sin en operand con S. Para ello, dispersa S con el mapa scatter_dims_to_operand_dims. De manera más formal:
   1. Sin[scatter_dims_to_operand_dims[k]] = S[k] si k < scatter_dims_to_operand_dims.size.
   2. Sin[_] = 0 en caso contrario.
3. Crea un índice Win en cada array operands dispersando los índices en update_window_dims en U según inserted_window_dims. De forma más formal:
   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 dominio [0, update_window_dims.size) y 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}, entonces window_dims_to_operand_dims es {0→1, 1→3, 2→4, 3→5}).
   2. Win[_] = 0 en caso contrario.
4. I es Win + Sin, en el que + es la adición por elemento.

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

• Inicializa output con operands, es decir, para todos los índices J, para todos los índices O en el array operands[J]:
  output[J][O] = operands[J][O]
• Para cada índice U en el 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 que se aplican las actualizaciones no es determinista. Por lo tanto, cuando varios índices en updates hagan 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 parámetro 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 establece como verdadero, XLA puede suponer que el usuario ordenó scatter_indices (en orden ascendente, después de dispersar sus valores según scatter_dims_to_operand_dims). Si no es así, la semántica se define en la implementación.

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

De manera informal, la operación de dispersión se puede ver como una operación inversa de la operación de recopilación, es decir, la operación de dispersión actualiza los elementos de la entrada que extrae la operación 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 array de salida a partir de elementos de dos arrays de entrada, en función de los valores de un array de predicado.

Select(pred, on_true, on_false)

Argumentos	Tipo	Semántica
pred	XlaOp	array of type 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 selecciones entre tuplas. Para este propósito, las tuplas se consideran 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 se puede considerar 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 a los índices de los elementos seleccionados para construir un array de salida con la misma forma que el array del operando. La función binaria select se usa para seleccionar un elemento de cada ventana aplicándolo en todas ellas, 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 muestra false si se selecciona el segundo parámetro, y la función debe mantener la 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 para una ventana determinada.

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

1. Es el valor actual en el índice seleccionado del 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 array de salida. Al principio, todos los índices del array de salida se establecen 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. SelectAndScatter se puede usar para propagar previamente los valores del 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	Es un 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 de cada ventana. Muestra true si se selecciona el primer parámetro y false si se selecciona el segundo.
window_dimensions	ArraySlice<int64>	Es un array de números enteros para los valores de la dimensión de ventana.
window_strides	ArraySlice<int64>	array de números enteros para valores de zancada de la ventana
padding	Padding	tipo de padding para la ventana (Padding::kSame o Padding::kValid)
source	XlaOp	array de tipo T con los valores que se dispersarán
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 origen de dispersión con su elemento de destino

En la siguiente imagen, 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) que aparece a continuación, un índice del array operand puede seleccionarse varias veces por diferentes ventanas. En la figura, ambas ventanas superiores (azul y roja) seleccionan el elemento de valor 9, y la función scatter de adición binaria produce el elemento de salida de valor 8 (2 + 6).

El orden de evaluación de la función scatter es arbitrario y puede no ser determinista. Por lo tanto, la función scatter no debe ser demasiado sensible a la reasociación. Consulta el análisis 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 (array de tipo T)
channel_handle	ChannelHandle	identificador único para cada par de envío/recepción

Envía los datos de 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 cliente de la operación Send representa la comunicación síncrona y se descompone internamente en 2 instrucciones de 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 usa una instrucción SendDone posterior para esperar a que se complete la transferencia de datos. El contexto es una tupla de {operando (forma), identificador de solicitud (U32)} y solo lo puede usar una instrucción SendDone.

SendDone(HloInstruction context)

Dado un contexto creado por una instrucción Send, espera a que se complete la transferencia de datos. La 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 de RecvDone
• Recv ocurre antes del RecvDone
• Send ocurre antes del SendDone

Cuando los compiladores de backend generan un programa lineal para cada cálculo que se comunica a través de instrucciones de canal, no debe haber ciclos en los cálculos. Por ejemplo, los programas que se muestran a continuación generan interbloqueos.

Porción

Consulta también XlaBuilder::Slice.

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

Slice(operand, start_indices, limit_indices, strides)

Argumentos	Tipo	Semántica
operand	XlaOp	Array de N dimensiones de tipo T
start_indices	ArraySlice<int64>	Es una lista de N números enteros que contienen los índices de inicio de la porción para cada dimensión. Los valores deben ser mayores o iguales que cero.
limit_indices	ArraySlice<int64>	Es una lista de N números enteros que contienen los índices de finalización (exclusivos) para el corte de cada dimensión. Cada valor debe ser superior o igual al valor start_indices correspondiente de la dimensión y menor o igual que el tamaño de la dimensión.
strides	ArraySlice<int64>	Es una lista de N números enteros que deciden el paso de entrada de la porción. El corte elige cada elemento strides[d] en la dimensión d.

Ejemplo de 1 dimensión:

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

Ejemplo de 2 dimensiones:

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 ordenarán.
comparator	XlaComputation	El cálculo del comparador que se usará.
dimension	int64	Es la dimensión en función de la cual se ordenan los datos.
is_stable	bool	Indica si se debe usar el ordenamiento estable.

Si solo se proporciona un operando:

• Si el operando es un tensor de rango-1 (una matriz), el resultado es una matriz ordenada. Si deseas ordenar el array en orden ascendente, el comparador debe realizar una comparación menor que. De forma formal, después de que se ordena el array, se mantiene para todas las posiciones de índice i, j con i < j que es 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 junto con la dimensión proporcionada. Por ejemplo, para un tensor de clasificación 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 manera independiente. Si no se proporciona un número de dimensión, se elige la última de forma predeterminada. Para la dimensión que se ordena, se aplica el mismo orden de clasificación que en el caso de la clasificación de rango 1.

Si se proporcionan operandos n > 1, ocurre lo siguiente:

• 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 forma 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 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 desempates.

• El resultado es una tupla que consta de los operandos en orden ordenado (junto con la dimensión proporcionada, como se indicó anteriormente). El operando i-th de la tupla corresponde al operando i-th de Ordenar.

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

Si is_stable se establece en verdadero, se garantiza que la ordenación sea estable, es decir, si el comparador considera que hay elementos 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. De forma predeterminada, is_stable se establece como falso.

Transposición

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

Transpose(operand)

Argumentos	Tipo	Semántica
operand	XlaOp	Es el operando que se transpone.
permutation	ArraySlice<int64>	Cómo permutar las dimensiones

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

Esto es lo mismo que Reshape(operand, permutation, Permute(permutation, operand.shape.dimensions)).

TriangularSolve

Consulta también XlaBuilder::TriangularSolve.

Resuelve sistemas de ecuaciones lineales con matrices de coeficientes triangulares inferiores o superiores mediante sustitución directa o inversa. Con la transmisión a lo largo de las dimensiones principales, esta rutina resuelve uno de los sistemas de matrices op(a) * x = b o x * op(a) = b para la variable x, dados a y b, donde 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 complejo o de punto flotante con forma [..., M, M].
b	XlaOp	Un array de rango superior a 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	Ya sea que uses a tal como está, lo transpongas o tomes 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 del otro triángulo están definidos por la implementación y pueden ser cualquier cosa.

Si el rango de a y b es superior a 2, se tratan como lotes de matrices, en los que 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 las tuplas (se puede acceder a ellas) a través de la operación GetTupleElement.

Mientras

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	Es el valor inicial del parámetro de condition y body.

Ejecuta body de forma secuencial hasta que falle condition. Esto es similar a un bucle while típico en muchos otros lenguajes, 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 nuevo resultado de 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 el pseudocódigo simplificado con un gráfico que representa el procesamiento. 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 sigue agregando 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};
}