Semántica de operaciones

bookmark_border Organiza tus páginas con colecciones Guarda y categoriza el contenido según tus preferencias.

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

Nota sobre la nomenclatura: el tipo de datos generalizado con el que se ocupa 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 se pueden enlazar entre operaciones con efectos secundarios 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)

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)

• 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)

• 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 alimentan hacen que el bucle 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 bloques a lo largo de split_dimensions, y los bloques se dispersan en todos los núcleos, p.ej., el bloque i se envía al núcleo i.
2. La fase de recopilación. Cada núcleo concatena los bloques recibidos a lo largo de concat_dimension.

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

• replica_groups: Cada ReplicaGroup contiene una lista de IDs de réplicas 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)

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 1, 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 a lo largo de la dimensión 0, 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)

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 los valores de los momentos en las dimensiones espaciales y de lotes.

El tipo de salida es una tupla de tres controladores:

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)

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

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

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)

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 de componentes en operand.

El algoritmo es el siguiente 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 epsilon, por lo general, es un número pequeño que se agrega para evitar errores de división por cero.

El tipo de salida es una tupla de tres XlaOp:

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)

Las dimensiones del operando y la forma objetivo deben coincidir, excepto 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 la operación completa 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)

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 la cantidad de dimensiones de un array duplicando los datos en él.

BroadcastInDim(operand, out_dim_size, broadcast_dimensions)

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

La aridad y los tipos de args deben coincidir con los parámetros de computation. Se permite no tener args.

CompositeCall

Consulta también XlaBuilder::CompositeCall.

Encapsula una operación compuesta por otras operaciones de StableHLO, toma entradas y atributos compuestos y produce resultados. El atributo de descomposición implementa la semántica de la operación. La operación compuesta se puede reemplazar por su descomposición sin cambiar la semántica del programa. En los casos en los que la incorporación de la descomposición no proporciona la misma semántica de operación, prefiere usar custom_call.

El campo de versión (que se establece de forma predeterminada en 0) se usa para indicar cuándo cambia la semántica de un compuesto.

Esta operación se implementa como un kCall con el atributo is_composite=true. El atributo computation especifica el campo decomposition. Los atributos del frontend almacenan los atributos restantes con el prefijo composite..

Ejemplo de operación CompositeCall:

f32[] call(f32[] %cst), to_apply=%computation, is_composite=true,
frontend_attributes = {
  composite.name="foo.bar",
  composite.attributes={n = 1 : i32, tensor = dense<1> : tensor<i32>},
  composite.version="1"
}

Call(computation, args..., name, composite_attributes, version)

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)

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 según la implementación y pueden ser cualquier cosa.

Si a tiene más de 2 dimensiones, a se considera un lote de matrices, en el que todas, excepto las 2 dimensiones menores, son dimensiones de lote.

Si a no es simétrica (hermítica) y definida positiva, el resultado se define según la implementación.

Restringir

Consulta también XlaBuilder::Clamp.

Limita un operando dentro del rango entre un valor mínimo y máximo.

Clamp(min, operand, max)

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 forma restringida de transmisión, min o max pueden ser un escalar de tipo T.

Ejemplo con min y max escalares:

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

Contraer

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

Contrae las dimensiones de un array en una.

Collapse(operand, dimensions)

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 sola dimensión nueva, en la misma posición de la secuencia de dimensiones que reemplazan, con el tamaño de la dimensión nueva igual al producto de los tamaños de las dimensiones originales. El número de dimensión más bajo en dimensions es la dimensión de variación más lenta (más importante) 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 datos entre réplicas.

CollectivePermute(operand, source_target_pairs)

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 de 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 la misma cantidad de dimensiones que cada uno de los operandos del array de entrada (que deben tener la misma cantidad de dimensiones entre sí) y contiene los argumentos en el orden en que se especificaron.

Concatenate(operands..., dimension)

Con la excepción de dimension, todas las dimensiones deben ser iguales. Esto se debe a que XLA no admite arrays “desordenados”. También ten en cuenta que los valores de 0 dimensiones 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)

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

true_computation debe aceptar un solo argumento del tipo $T_0$ y se invocará con true_operand, que debe ser del mismo tipo. false_computation debe aceptar un solo argumento del 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)

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. VALID padding 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í, una convolución se puede considerar como una ventana n-dimensional que se mueve a través de un área base n-dimensional, y se realiza un cálculo para cada posición posible de la ventana.

Supongamos que n es la cantidad de dimensiones espaciales. El argumento lhs es un array de (n+2) dimensiones que describe el área base. Esto se denomina entrada, aunque, por supuesto, la parte derecha 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 de 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 (n+2) dimensiones que describe el filtro, el kernel o la ventana de convolución. 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 de 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 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 orificios se completan con un valor de no operación, 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 atributos de entrada y 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 i se usan juntos para calcular feature_group_count para muchas convoluciónes 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 filtros agrupados durante la retropropagación. batch_group_count debe ser un divisor del tamaño de la dimensión del lote lhs (entrada). Si batch_group_count es mayor que 1, significa que la dimensión del lote de salida debe tener el tamaño 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. En efecto, dividimos cada lote de lhs en grupos batch_group_count y hacemos lo mismo con las funciones de salida. Luego, para cada uno de estos grupos, realizamos contracciones por pares y concatenamos el resultado a lo largo de la dimensión de atributos de salida. La semántica operativa de todas las demás dimensiones (espacial 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 una posición de una esquina de la ventana dentro del área de 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. Los dos cuadros tienen las mismas dimensiones, por lo que podemos tomar la suma de los productos por elemento entre los dos cuadros (similar a un producto punto). 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)

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

Una conversión como T=s32 a U=f32 realizará una rutina de conversión de int a flotante normalizada, como redondear al número 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)

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)

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

La operación realiza la suma de productos en la segunda dimensión de lhs (o la primera si tiene 1 dimensión) 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)

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

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

No es necesario que los números de dimensión de contratación asociados de lhs y rhs sean 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 de lote asociados de lhs y rhs deben tener los mismos tamaños de dimensión.

Ejemplo con números de dimensión del lote (tamaño del lote 2, matrices 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} } }

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 1 dimensión, con un tamaño de dimensión igual a la cantidad de dimensiones de operand.

DynamicSlice(operand, start_indices, size_indices)

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 que se reemplaza en start_indices. La forma de update determina la forma del subarray del resultado que se actualiza. La forma de start_indices debe ser unidimensional, con un tamaño de dimensión igual a la cantidad de dimensiones de operand.

DynamicUpdateSlice(operand, update, start_indices)

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: 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 (arctg con 2 argumentos) o Complex (combina partes reales e imaginarias en un número complejo)

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 arreglos 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/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 con transmisiones de diferentes dimensiones para estas operaciones:

Op(lhs, rhs, broadcast_dimensions)

En el que Op es igual que el anterior. Esta variante de la operación debe usarse para operaciones aritméticas 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 se usa para expandir la cantidad de dimensiones del operando de dimensión inferior hasta la cantidad de dimensiones del operando de dimensión superior. broadcast_dimensions asigna las dimensiones de la forma de menor dimensión a las dimensiones de la forma de mayor dimensión. Las dimensiones no asignadas de la forma expandida se completan con dimensiones de tamaño uno. Luego, la transmisión de dimensiones degeneradas 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 se aplica la semántica de comparación de punto flotante estándar IEEE 754 cuando se comparan tipos de punto flotante.

Op(lhs, rhs)

En el que Op es uno de 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.

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 dimensiones para estas operaciones:

Op(lhs, rhs, broadcast_dimensions)

En el que Op es igual que el anterior. Esta variante de la operación debe usarse 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 por 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) Exponente natural por elemento x -> e^x.

Expm1(operand) Exponencial natural por elemento menos unox -> 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 por 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 elemento x -> -x.

Not(operand) No lógico de elemento inteligente x -> !(x).

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

Real(operand) Es la parte real de un elemento 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 raíz cuadrada por elemento x -> sqrt(x).

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

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

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 (0 dimensiones).

Fft

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

Consulta también XlaBuilder::Fft.

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

La FFT de CPU está respaldada por 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)

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 con batch_dims.size + offset_dims.size dimensiones.

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 esta función está bien definida incluso si G está vacía: si G está vacía, S = start_indices.

2. Crea un índice inicial, Sin, en operand con S. Para ello, dispérsalo con start_index_map.S 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, en el que + es la adición por elemento.

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 ordenó 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.

• Dividimos de forma dinámica una porción con el tamaño slice_sizes usando el índice de partida 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, entonces, 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 que se describió anteriormente 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 compensación de salida presentes de forma explícita en la forma de salida puede ser menor que la cantidad de dimensiones de entrada. Estas dimensiones "faltantes", que se enumeran de forma explícita 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 omitirlos no introduce ambigüedades.

3. La porción extraída del array "Gather Indices" ((X, Y) en el último ejemplo) puede tener menos elementos que la cantidad de dimensiones del array de entrada, y una asignación explícita dicta cómo se debe expandir el índice para tener la misma cantidad de dimensiones que la entrada.

Como ejemplo final, 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)

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)

Pasa el operando como resultado, con la dimensión dinámica a la que hace un seguimiento el compilador.

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 del tiempo de compilación para que la inferencia de forma pueda determinar el tipo del valor resultante.

Esto es similar 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)

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 feeds in-feed en el 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 no válida y continúa sin leer ningún dato del Infeed del dispositivo.

Pizca

Consulta también XlaBuilder::Iota.

Iota(shape, iota_dimension)

Compila una constante literal 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.

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)

Aplica una función escalar a 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)

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 interior ocurre, de forma lógica, antes del padding de borde, por lo que, en el caso del padding de borde negativo, se quitan elementos 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)

Recibe datos de la forma determinada de una instrucción Send en otra computación que comparte el mismo identificador 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 descompone internamente en 2 instrucciones de 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 usa una instrucción RecvDone posterior 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)

Donde:

• 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. La cantidad de dimensiones de cada array que se muestra es number_of_dimensions(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 números de punto flotante es lo suficientemente cercana como para ser asociativa para la mayoría de los usos prácticos.

Ejemplos

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

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 number of dimensions 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 2 dimensiones: la dimensión 0 de tamaño 2 y la 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 cantidad de dimensiones es 3, la dimensión 0 de tamaño 4, la dimensión 1 de tamaño 2 y la dimensión 2 de tamaño 3. Para simplificar, los valores del 1 al 6 se replican en 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 2 dimensiones 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 2 dimensiones 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 cantidad de dimensiones).

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 del 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 de exponente y mantisa en el formato de menor precisión se puede especificar de forma arbitraria, aunque es posible que no todos los tamaños de bits sean compatibles con todas las implementaciones de hardware.

ReducePrecision(operand, mantissa_bits, exponent_bits)

El resultado es un array de tipo T. Los valores de entrada se redondean al valor más cercano representable con la cantidad determinada de bits de mantisa (con la semántica de "empates en par") y cualquier valor que exceda el rango especificado por la cantidad de bits de exponente se limita 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 de exponente (para distinguir un valor cero de un infinito, ya que ambos tienen una mantisa cero) y debe tener una cantidad no negativa de bits de mantisa. La cantidad de bits de exponente o mantisa puede exceder el valor correspondiente para el tipo T. La parte correspondiente de la conversión es, entonces, simplemente una operación no realizada.

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)

• 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 grupos de réplicas entre los que se realiza la reducción (el ID de réplica de la réplica actual se puede recuperar con ReplicaId). El orden de las réplicas en cada grupo determina el orden en el que se dispersará el resultado de la reducción total. 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 tener el mismo tamaño. Por ejemplo, replica_groups = {0, 2}, {1, 3} realiza la reducción entre las réplicas 0 y 2, y 1 y 3, y luego dispersa el resultado.
• shard_count es el tamaño de cada grupo de réplicas. Necesitamos esto en los casos en 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 de cada ventana de una secuencia de N آرایه‌های چند بعدی، lo que produce una sola tupla o N آرایه‌های چند بعدی 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)

Donde:

• 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, las dimensiones de ventana_paso deben ser iguales a las dimensiones de ventana. En la siguiente figura, se ilustra el uso de dos valores de paso 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 y obteniendo [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.

Devuelve 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, dimensions)

Conceptualmente, reshape primero aplana un array en un vector unidimensional de valores de datos y, luego, define mejor este vector en una forma nueva. 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. El vector dimensions determina el tamaño del array de salida. El valor en el índice 0 en dimensions 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 de dimensions 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 dimensions, las dimensiones en dimensions se ordenan de la variación más lenta (más principal) a la más rápida (más menor).

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} } };

let v012_24 = Reshape(v, {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, {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} };

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

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

Rev (invertido)

Consulta también XlaBuilder::Rev.

Rev(operand, dimensions)

Invierte el orden de los elementos del array operand a lo largo del dimensions especificado y genera un array de salida del mismo formato. 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 valores escalares.

RngNormal(mu, sigma, shape)

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)

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

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)

Valores disponibles para algorithm:

• rng_default: Algoritmo específico del backend con requisitos de forma específicos del 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)

Donde:

• N debe ser 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 tiene una dimensión 1 final.

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

Los argumentos de dispersión deben seguir estas restricciones:

• Cada array updates debe tener update_window_dims.size + scatter_indices.rank - 1 dimensiones.

• 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, es igual a update_window_dims[k] para algún k), el límite de la dimensión i en updates no debe exceder el límite correspondiente de operand después de tener en cuenta inserted_window_dims (es decir, adjusted_window_bounds[k], donde adjusted_window_bounds contiene los límites de operand con los límites en los índices inserted_window_dims quitados).
  • 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 debe tener números de dimensión repetidos y debe 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 forma 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 a los que se dispersan no son únicos, la semántica se define en la implementación.

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, según los valores de un array de predicados.

Select(pred, on_true, on_false)

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 select binaria se usa para seleccionar un elemento de cada ventana aplicándolo en cada ventana, y se lo llama con la propiedad de que el vector de índices del primer parámetro es lexicográficamente menor que el vector de índices 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. Inicialmente, 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 propagarse hacia atrás los valores del gradiente de una capa de agregación en una red neuronal.

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

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, las dos 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 ser no 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)

Envía los datos del operando determinado a una instrucción Recv en otro procesamiento que comparte el mismo identificador 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 que se complete la transferencia de datos. El contexto es una tupla de {operando (forma), identificador de solicitud (U32)} y solo puede ser utilizada por 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 de Send
• Send ocurre antes de RecvDone
• Recv ocurre antes de RecvDone
• Send ocurre antes de 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 indican a continuación generan interbloqueos.

Ten en cuenta que la restricción de las instrucciones solo se aplica a las TPU en el tiempo de ejecución. En la GPU, send y recv bloquearán y no enviarán datos reales hasta después de un protocolo de enlace entre los dispositivos de origen y destino.

Porción

Consulta también XlaBuilder::Slice.

El corte extrae un subarray del array de entrada. El subarray tiene la misma cantidad de dimensiones que la entrada y contiene los valores dentro de un cuadro delimitador dentro del array de entrada, donde las dimensiones y los índices del cuadro delimitador se proporcionan como argumentos a la operación de corte.

Slice(operand, start_indices, limit_indices, strides)

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)

Si solo se proporciona un operando, ocurre lo siguiente:

• Si el operando es un tensor de 1 dimensión (un array), el resultado es un array ordenado. 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 sean comparator(value[i], value[j]) = comparator(value[j], value[i]) = false o comparator(value[i], value[j]) = true.

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

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

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

TopK

Consulta también XlaBuilder::TopK.

TopK encuentra los valores y los índices de los elementos más grandes o más pequeños de k para la última dimensión del tensor determinado.

TopK(operand, k, largest)

Para un tensor de entrada de 1 dimensión (un array), encuentra las entradas k más grandes o más pequeñas del array y muestra una tupla de dos arrays (values, indices). Por lo tanto, values[j] es la entrada j más grande o más pequeña en operand, y su índice es indices[j].

Para un tensor de entrada con más de 1 dimensión, calcula las k entradas principales a lo largo de la última dimensión y conserva todas las demás dimensiones (filas) en el resultado. Por lo tanto, para un operando de forma [A, B, ..., P, Q] en el que Q >= k, el resultado es una tupla (values, indices) en la que se cumple lo siguiente:

values.shape = indices.shape = [A, B, ..., P, k]

Si dos elementos de una fila son iguales, el elemento con el índice más bajo aparece primero.

Transposición

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

Transpose(operand)

Permuta las dimensiones del operando con la permutación dada, por lo que es ∀ i . 0 ≤ i < number of dimensions ⇒ 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. A través de 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)