Semántica de operaciones

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

Una nota sobre la nomenclatura: El tipo de datos generalizado con el que trabaja XLA es un array N-dimensional 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 unidimensional y una matriz es un array bidimensional.

Obtén más información sobre la estructura de una operación en Formas y diseño y Diseño en mosaicos.

Abdominales

Consulta también XlaBuilder::Abs.

Valor absoluto de cada elemento de x -> |x|.

Abs(operand)

Para obtener información sobre StableHLO, consulta StableHLO: abs.

Agregar

Consulta también XlaBuilder::Add.

Realiza la suma de lhs y rhs elemento por elemento.

Add(lhs, rhs)

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

Existe una variante alternativa con compatibilidad de transmisión multidimensional para Add:

Add(lhs,rhs, broadcast_dimensions)

Esta variante de la operación se debe usar para operaciones aritméticas entre arrays de diferentes rangos (como sumar una matriz a un vector).

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

Para obtener información sobre StableHLO, consulta StableHLO: add.

AddDependency

Consulta también HloInstruction::AddDependency.

AddDependency puede aparecer en volcados de HLO, pero no está diseñado para que los usuarios finales lo construyan manualmente.

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 encadenar entre operaciones con efectos secundarios para aplicar un orden. AfterAll se puede usar como una unión de tokens para ordenar una operación después de un conjunto de operaciones.

AfterAll(tokens)

Para obtener información sobre StableHLO, consulta StableHLO - after_all.

AllGather

Consulta también XlaBuilder::AllGather.

Realiza la concatenación en las réplicas.

AllGather(operand, all_gather_dimension, shard_count, replica_groups, channel_id, layout, use_global_device_ids)

• replica_groups es una lista de grupos de réplicas entre los que se realiza la concatenación (el ID de la réplica actual se puede recuperar con ReplicaId). El orden de las réplicas en cada grupo determina el orden en el que se 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 este campo en los casos en que replica_groups esté vacío.
• channel_id se usa para la comunicación entre módulos: solo las operaciones de all-gather con el mismo channel_id pueden comunicarse entre sí.
• use_global_device_ids Devuelve verdadero si los IDs de la configuración de ReplicaGroup representan un ID global de (replica_id * partition_count + partition_id) en lugar de un ID de réplica. Esto permite una agrupación más flexible de los dispositivos si esta operación de reducción total es tanto entre particiones como entre réplicas.

La forma de salida es la forma de entrada con el all_gather_dimension aumentado shard_count veces. Por ejemplo, si hay dos réplicas y el operando tiene el valor [1.0, 2.5] y [3.0, 5.25], respectivamente, en las dos réplicas, el valor de salida de esta operación en la que all_gather_dim es 0 será [1.0, 2.5, 3.0,5.25] en ambas réplicas.

Internamente, la API de AllGather se descompone en 2 instrucciones de HLO (AllGatherStart y AllGatherDone).

Consulta también HloInstruction::CreateAllGatherStart.

AllGatherStart y AllGatherDone sirven como primitivas en HLO. Estas operaciones pueden aparecer en volcados de HLO, pero no están diseñadas para que los usuarios finales las construyan manualmente.

Para obtener información sobre StableHLO, consulta StableHLO - all_gather.

AllReduce

Consulta también XlaBuilder::AllReduce.

Realiza un cálculo personalizado en todas las réplicas.

AllReduce(operand, computation, replica_groups, channel_id, shape_with_layout, use_global_device_ids)

• Cuando operand es una tupla de arrays, la reducción de todos 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 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 de all-reduce con el mismo channel_id pueden comunicarse entre sí.
• shape_with_layout: Fuerza el diseño de AllReduce al diseño determinado. Se usa para garantizar el mismo diseño para un grupo de operaciones de AllReduce compiladas por separado.
• use_global_device_ids Devuelve verdadero si los IDs de la configuración de ReplicaGroup representan un ID global de (replica_id * partition_count + partition_id) en lugar de un ID de réplica. Esto permite una agrupación más flexible de los dispositivos si esta operación de reducción total es tanto entre particiones como entre réplicas.

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

Para calcular el resultado de AllReduce, se requiere una entrada de cada réplica, por lo que, si una réplica ejecuta un nodo AllReduce más veces que otra, la primera esperará indefinidamente. Dado que todas las réplicas ejecutan el mismo programa, no hay muchas formas en que esto pueda suceder, pero es posible cuando la condición de un bucle while depende de los datos de infeed y los datos que son infeed hacen que el bucle while itere más veces en una réplica que en otra.

Internamente, la API de AllReduce se descompone en 2 instrucciones de HLO (AllReduceStart y AllReduceDone).

Consulta también HloInstruction::CreateAllReduceStart.

AllReduceStart y AllReduceDone sirven como primitivas en HLO. Estas operaciones pueden aparecer en volcados de HLO, pero no están diseñadas para que los usuarios finales las construyan manualmente.

CrossReplicaSum

Consulta también XlaBuilder::CrossReplicaSum.

Realiza AllReduce con un cálculo de suma.

CrossReplicaSum(operand, replica_groups)

Devuelve la suma del valor del operando dentro de cada subgrupo de réplicas. Todas las réplicas proporcionan una entrada a la suma y todas las réplicas reciben la suma resultante para cada subgrupo.

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-ésimo se envía al núcleo i-ésimo.
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 las siguientes maneras:

• replica_groups: Cada ReplicaGroup contiene una lista de IDs de réplica que participan en el cálculo (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 y 3. Luego, se aplicará otro 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, layout, channel_id)

Consulta xla::shapes para obtener más información sobre las formas y los diseños.

Para obtener información sobre StableHLO, consulta StableHLO: all_to_all.

AllToAll: ejemplo 1

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

AllToAll, ejemplo 2, StableHLO

En el ejemplo anterior, hay 2 réplicas que participan en AllToAll. En cada réplica, el operando tiene la forma f32[2,4]. El operando se divide en 2 partes a lo largo de la dimensión 1, por lo que cada parte tiene la forma f32[2,2]. Luego, las 2 partes se intercambian entre las réplicas según su posición en el grupo de réplicas. Cada réplica recopila su parte correspondiente de ambos operandos y los concatena a lo largo de la dimensión 0. Como resultado, la salida en cada réplica tiene la forma f32[4,2].

RaggedAllToAll

Consulta también XlaBuilder::RaggedAllToAll.

RaggedAllToAll realiza una operación colectiva de todos a todos, en la que la entrada y la salida son tensores irregulares.

RaggedAllToAll(input, input_offsets, send_sizes, output, output_offsets, recv_sizes, replica_groups, channel_id)

Los tensores irregulares se definen con un conjunto de tres tensores:

• data: El tensor data es “irregular” a lo largo de su dimensión más externa, a lo largo de la cual cada elemento indexado tiene un tamaño variable.
• offsets: El tensor offsets indexa la dimensión más externa del tensor data y representa el desplazamiento inicial de cada elemento irregular del tensor data.
• sizes: El tensor sizes representa el tamaño de cada elemento irregular del tensor data, en el que el tamaño se especifica en unidades de subelementos. Un subelemento se define como el sufijo de la forma del tensor "data" que se obtiene quitando la dimensión "irregular" más externa.
• Los tensores offsets y sizes deben tener el mismo tamaño.

Ejemplo de un tensor irregular:

data: [8,3] =
{ {a,b,c},{d,e,f},{g,h,i},{j,k,l},{m,n,o},{p,q,r},{s,t,u},{v,w,x} }

offsets: [3] = {0, 1, 4}

sizes: [3] = {1, 3, 4}

// Index 'data' at 'offsets'[0], 'sizes'[0]' // {a,b,c}

// Index 'data' at 'offsets'[1], 'sizes'[1]' // {d,e,f},{g,h,i},{j,k,l}

// Index 'data' at 'offsets'[2], 'sizes'[2]' // {m,n,o},{p,q,r},{s,t,u},{v,w,x}

output_offsets debe fragmentarse de manera que cada réplica tenga desplazamientos en la perspectiva de salida de la réplica de destino.

Para el desplazamiento de salida i-ésimo, la réplica actual enviará la actualización input[input_offsets[i]:input_offsets[i]+input_sizes[i]] a la réplica i-ésima que se escribirá en output_i[output_offsets[i]:output_offsets[i]+send_sizes[i]] en la réplica i-ésima output.

Por ejemplo, si tenemos 2 réplicas:

replica 0:
input: [1, 2, 2]
output:[0, 0, 0, 0]
input_offsets: [0, 1]
send_sizes: [1, 2]
output_offsets: [0, 0]
recv_sizes: [1, 1]

replica 1:
input: [3, 4, 0]
output: [0, 0, 0, 0]
input_offsets: [0, 1]
send_sizes: [1, 1]
output_offsets: [1, 2]
recv_sizes: [2, 1]

// replica 0's result will be: [1, 3, 0, 0]
// replica 1's result will be: [2, 2, 4, 0]

El HLO de todos a todos irregular tiene los siguientes argumentos:

• input: Tensor de datos de entrada irregulares.
• output: Tensor de datos de salida irregulares.
• input_offsets: Tensor de desplazamiento de entrada irregular.
• send_sizes: Tensor de tamaños de envío irregulares.
• output_offsets: Es un array de desplazamientos irregulares en el resultado de la réplica de destino.
• recv_sizes: Es un tensor de tamaños de recv irregulares.

Todos los tensores *_offsets y *_sizes deben tener la misma forma.

Se admiten dos formas para los tensores *_offsets y *_sizes:

• [num_devices], donde ragged-all-to-all puede enviar como máximo una actualización a cada dispositivo remoto del grupo de réplicas. Por ejemplo:

for (remote_device_id : replica_group) {
     SEND input[input_offsets[remote_device_id]],
     output[output_offsets[remote_device_id]],
     send_sizes[remote_device_id] }

• [num_devices, num_updates], donde ragged-all-to-all puede enviar hasta num_updates actualizaciones al mismo dispositivo remoto (cada una con diferentes compensaciones), para cada dispositivo remoto del grupo de réplicas.

Por ejemplo:

for (remote_device_id : replica_group) {
    for (update_idx : num_updates) {
        SEND input[input_offsets[remote_device_id][update_idx]],
        output[output_offsets[remote_device_id][update_idx]]],
        send_sizes[remote_device_id][update_idx] } }

y

Consulta también XlaBuilder::And.

Realiza un AND a nivel de los elementos de dos tensores lhs y rhs.

And(lhs, rhs)

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

Existe una variante alternativa con compatibilidad para la transmisión de diferentes dimensiones para And:

And(lhs,rhs, broadcast_dimensions)

Esta variante de la operación se debe usar para operaciones aritméticas entre arrays de diferentes rangos (como sumar una matriz a un vector).

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

Para obtener información sobre StableHLO, consulta StableHLO: y.

Asíncrono

Consulta también HloInstruction::CreateAsyncStart, HloInstruction::CreateAsyncUpdate, HloInstruction::CreateAsyncDone.

AsyncDone, AsyncStart y AsyncUpdate son instrucciones internas de HLO que se usan para las operaciones asíncronas y sirven como primitivas en HLO. Estas operaciones pueden aparecer en volcados de HLO, pero no están diseñadas para que los usuarios finales las construyan manualmente.

Atan2

Consulta también XlaBuilder::Atan2.

Realiza la operación atan2 en términos de elementos en lhs y rhs.

Atan2(lhs, rhs)

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

Existe una variante alternativa con compatibilidad para la transmisión de diferentes dimensiones para Atan2:

Atan2(lhs,rhs, broadcast_dimensions)

Esta variante de la operación se debe usar para operaciones aritméticas entre arrays de diferentes rangos (como sumar una matriz a un vector).

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

Para obtener información sobre StableHLO, consulta StableHLO: atan2.

BatchNormGrad

Consulta también XlaBuilder::BatchNormGrad y el informe 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, batch_mean, batch_var, grad_output, 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 los gradientes con respecto a operand, offset y scale en todas las demás dimensiones. El 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 la característica 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 batch_mean y batch_var representan valores de momentos en las dimensiones espaciales y de lote.

El tipo de salida es una tupla de tres identificadores:

Para obtener información sobre StableHLO, consulta StableHLO - batch_norm_grad.

BatchNormInference

Consulta también XlaBuilder::BatchNormInference y el informe 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. El feature_index debe ser un índice válido para la dimensión de la característica 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.

Para obtener información sobre StableHLO, consulta StableHLO: batch_norm_inference.

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. El feature_index debe ser un índice válido para la dimensión de la característica en operand.

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

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

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

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

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

El tipo de salida es una tupla de tres XlaOp:

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

Para obtener información sobre StableHLO, consulta StableHLO - batch_norm_training.

Bitcast

Consulta también HloInstruction::CreateBitcast.

Bitcast puede aparecer en los volcados de HLO, pero no está diseñado para que los usuarios finales lo construyan manualmente.

BitcastConvertType

Consulta también XlaBuilder::BitcastConvertType.

De manera similar a un tf.bitcast en TensorFlow, realiza una operación de bitcast elemento por elemento desde una forma de datos a una forma de destino. 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 conversión de bits, y un elemento s32 se convertirá en cuatro elementos s8. Bitcast se implementa como una conversión de bajo nivel, por lo que las máquinas con diferentes representaciones de punto flotante arrojará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 primitivo antes y después de la conversión.

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

Para obtener información sobre StableHLO, consulta StableHLO - bitcast_convert.

Conversión de Bitcast a un tipo primitivo de ancho diferente

La instrucción BitcastConvert de HLO admite el caso en el que el tamaño del tipo de elemento de salida T' no es igual al tamaño del elemento de entrada T. Como toda la operación es, conceptualmente, una conversión de bits y no cambia los bytes subyacentes, la forma del elemento de salida debe cambiar. En el caso de 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 efectivos:

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

Transmisión

Consulta también XlaBuilder::Broadcast.

Agrega dimensiones a un array duplicando los datos en el array.

Broadcast(operand, broadcast_sizes)

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

Los nuevos índices de dimensiones se copian en los operandos, es decir:

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

Por ejemplo, si operand es un escalar f32 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.

Para obtener información sobre StableHLO, consulta StableHLO: broadcast.

BroadcastInDim

Consulta también XlaBuilder::BroadcastInDim.

Expande el tamaño y la cantidad de dimensiones de un array duplicando los datos en el array.

BroadcastInDim(operand, out_dim_size, broadcast_dimensions)

Es similar a Broadcast, pero permite agregar dimensiones en cualquier lugar y expandir las dimensiones existentes con tamaño 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 i-ésima dimensión del operando se asigna a la dimensión broadcast_dimension[i] de la forma de salida. Las dimensiones de operand deben tener un tamaño de 1 o 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 realiza a lo largo de estas dimensiones degeneradas para alcanzar la forma del resultado. 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, operands...)

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

CompositeCall

Consulta también XlaBuilder::CompositeCall.

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

El campo de versión (que tiene el valor predeterminado 0) se usa para indicar cuándo cambian las semánticas de un elemento compuesto.

Este op se implementa como un kCall con el atributo is_composite=true. El campo decomposition se especifica con el atributo computation. Los atributos de 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"
}

CompositeCall(computation, operands..., name, attributes, version)

El decomposition de una operación no es un campo llamado, sino que aparece como un atributo to_apply que apunta a la función que contiene la implementación de nivel inferior, es decir, to_apply=%funcname.

Puedes encontrar más información sobre la composición y la descomposición en la Especificación de StableHLO.

Cbrt

Consulta también XlaBuilder::Cbrt.

Operación de raíz cúbica para cada elemento x -> cbrt(x).

Cbrt(operand)

Cbrt también admite el argumento opcional result_accuracy:

Cbrt(operand, result_accuracy)

Para obtener más información sobre result_accuracy, consulta Precisión de los resultados.

Para obtener información sobre StableHLO, consulta StableHLO - cbrt.

Techo

Consulta también XlaBuilder::Ceil.

Función ceil para cada elemento x -> ⌈x⌉.

Ceil(operand)

Para obtener información sobre StableHLO, consulta StableHLO: ceil.

Cholesky

Consulta también XlaBuilder::Cholesky.

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

Cholesky(a, lower)

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

Los datos de entrada se leen solo desde el 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 devuelven en el mismo triángulo; los valores del otro triángulo se definen según la implementación y pueden ser cualquier valor.

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

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

Para obtener información sobre StableHLO, consulta StableHLO: cholesky.

Restringir

Consulta también XlaBuilder::Clamp.

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

Clamp(min, operand, max)

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

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

Ejemplo con min y max escalares:

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

Para obtener información sobre StableHLO, consulta StableHLO: clamp.

Contraer

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

Reduce las dimensiones de un array a una sola dimensión.

Collapse(operand, dimensions)

El 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 constante en tiempo de compilación de índices de dimensión. Los índices de dimensión deben ser un subconjunto consecutivo y ordenado (de menor a mayor) 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 las que reemplazan, con el tamaño de la nueva dimensión 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 (la más importante) en el bucle anidado que contrae estas dimensiones, y el número de dimensión más alto es la dimensión de variación más rápida (la menos importante). Consulta el operador tf.reshape si necesitas un orden de contracción más general.

Por ejemplo, sea v 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} };

Clz

Consulta también XlaBuilder::Clz.

Cuenta los ceros iniciales de cada elemento.

Clz(operand)

CollectiveBroadcast

Consulta también XlaBuilder::CollectiveBroadcast.

Difunde datos entre las réplicas. Los datos se envían desde el primer ID de réplica de cada grupo a los demás IDs del mismo grupo. Si un ID de réplica no está en ningún grupo de réplicas, el resultado en esa réplica es un tensor que consta de 0 en shape.

CollectiveBroadcast(operand, replica_groups, channel_id)

Para obtener información sobre StableHLO, consulta StableHLO - collective_broadcast.

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, channel_id, inplace)

Ten en cuenta que existen las siguientes restricciones en source_target_pairs:

• Ningún par debe tener el mismo ID de réplica de destino ni el mismo ID de réplica de origen.
• Si un ID de réplica no es un destino en ningún par, el resultado en esa réplica es un tensor que consta de 0(s) con la misma forma que la entrada.

Internamente, la API de la operación CollectivePermute se descompone en 2 instrucciones de HLO (CollectivePermuteStart y CollectivePermuteDone).

Consulta también HloInstruction::CreateCollectivePermuteStart.

CollectivePermuteStart y CollectivePermuteDone sirven como primitivas en HLO. Estas operaciones pueden aparecer en los volcados de HLO, pero no están diseñadas para que los usuarios finales las construyan manualmente.

Para obtener información sobre StableHLO, consulta StableHLO - collective_permute.

Comparar

Consulta también XlaBuilder::Compare.

Realiza una comparación elemento por elemento de lhs y rhs de lo siguiente:

Ec.

Consulta también XlaBuilder::Eq.

Realiza una comparación igual a elemento por elemento de lhs y rhs.

\(lhs = rhs\)

Eq(lhs, rhs)

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

Existe una variante alternativa con compatibilidad para la transmisión de diferentes dimensiones para Eq:

Eq(lhs,rhs, broadcast_dimensions)

Esta variante de la operación se debe usar para operaciones aritméticas entre arrays de diferentes rangos (como sumar una matriz a un vector).

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

Se admite un orden total sobre los números de punto flotante para Eq, ya que se aplica lo siguiente:

\[-NaN < -Inf < -Finite < -0 < +0 < +Finite < +Inf < +NaN.\]

EqTotalOrder(lhs,rhs, broadcast_dimensions)

Para obtener información sobre StableHLO, consulta StableHLO: Comparación.

Ne

Consulta también XlaBuilder::Ne.

Realiza una comparación de no igual a elemento por elemento de lhs y rhs.

\(lhs != rhs\)

Ne(lhs, rhs)

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

Existe una variante alternativa con compatibilidad para la transmisión de diferentes dimensiones para Ne:

Ne(lhs,rhs, broadcast_dimensions)

Esta variante de la operación se debe usar para operaciones aritméticas entre arrays de diferentes rangos (como sumar una matriz a un vector).

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

Existe un orden total sobre los números de punto flotante para Ne, ya que se aplica lo siguiente:

\[-NaN < -Inf < -Finite < -0 < +0 < +Finite < +Inf < +NaN.\]

NeTotalOrder(lhs,rhs, broadcast_dimensions)

Para obtener información sobre StableHLO, consulta StableHLO: Comparación.

Ge

Consulta también XlaBuilder::Ge.

Realiza una comparación greater-or-equal-than según cada elemento de lhs y rhs.

\(lhs >= rhs\)

Ge(lhs, rhs)

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

Existe una variante alternativa con compatibilidad para la transmisión de diferentes dimensiones para Ge:

Ge(lhs,rhs, broadcast_dimensions)

Esta variante de la operación se debe usar para operaciones aritméticas entre arrays de diferentes rangos (como sumar una matriz a un vector).

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

Existe un orden total sobre los números de punto flotante para Gt, ya que se aplica lo siguiente:

\[-NaN < -Inf < -Finite < -0 < +0 < +Finite < +Inf < +NaN.\]

GtTotalOrder(lhs,rhs, broadcast_dimensions)

Para obtener información sobre StableHLO, consulta StableHLO: Comparación.

Gt

Consulta también XlaBuilder::Gt.

Realiza una comparación mayor que a nivel del elemento de lhs y rhs.

\(lhs > rhs\)

Gt(lhs, rhs)

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

Existe una variante alternativa con compatibilidad para la transmisión de diferentes dimensiones para Gt:

Gt(lhs,rhs, broadcast_dimensions)

Esta variante de la operación se debe usar para operaciones aritméticas entre arrays de diferentes rangos (como sumar una matriz a un vector).

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

Para obtener información sobre StableHLO, consulta StableHLO: Comparación.

Le

Consulta también XlaBuilder::Le.

Realiza una comparación less-or-equal-than a nivel del elemento de lhs y rhs.

\(lhs <= rhs\)

Le(lhs, rhs)

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

Existe una variante alternativa con compatibilidad para la transmisión de diferentes dimensiones para Le:

Le(lhs,rhs, broadcast_dimensions)

Esta variante de la operación se debe usar para operaciones aritméticas entre arrays de diferentes rangos (como sumar una matriz a un vector).

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

Existe un orden total sobre los números de punto flotante para Le, ya que se aplica lo siguiente:

\[-NaN < -Inf < -Finite < -0 < +0 < +Finite < +Inf < +NaN.\]

LeTotalOrder(lhs,rhs, broadcast_dimensions)

Para obtener información sobre StableHLO, consulta StableHLO: Comparación.

Lt

Consulta también XlaBuilder::Lt.

Realiza una comparación menor que elemento por elemento de lhs y rhs.

\(lhs < rhs\)

Lt(lhs, rhs)

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

Existe una variante alternativa con compatibilidad para la transmisión de diferentes dimensiones para Lt:

Lt(lhs,rhs, broadcast_dimensions)

Esta variante de la operación se debe usar para operaciones aritméticas entre arrays de diferentes rangos (como sumar una matriz a un vector).

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

Existe un orden total sobre los números de punto flotante para Lt, ya que se aplica lo siguiente:

\[-NaN < -Inf < -Finite < -0 < +0 < +Finite < +Inf < +NaN.\]

LtTotalOrder(lhs,rhs, broadcast_dimensions)

Para obtener información sobre StableHLO, consulta StableHLO: Comparación.

Complejo

Consulta también XlaBuilder::Complex.

Realiza una conversión de cada elemento a un valor complejo a partir de un par de valores reales e imaginarios, lhs y rhs.

Complex(lhs, rhs)

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

Existe una variante alternativa con compatibilidad de transmisión multidimensional para Complex:

Complex(lhs,rhs, broadcast_dimensions)

Esta variante de la operación se debe usar para operaciones aritméticas entre arrays de diferentes rangos (como sumar una matriz a un vector).

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

Para obtener información sobre StableHLO, consulta StableHLO: complejo.

ConcatInDim (Concatenate)

Consulta también XlaBuilder::ConcatInDim.

La función 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 "irregulares". 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 unidimensional:

Concat({ {2, 3}, {4, 5}, {6, 7} }, 0)
//Output:  {2, 3, 4, 5, 6, 7}

Ejemplo bidimensional:

let a = { {1, 2},
         {3, 4},
         {5, 6} };

let b = { {7, 8} };

Concat({a, b}, 0)

//Output:  { {1, 2},
//          {3, 4},
//          {5, 6},
//          {7, 8} }

Diagrama:

Para obtener información sobre StableHLO, consulta StableHLO: concatenate.

Condicional

Consulta también XlaBuilder::Conditional.

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

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

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

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

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, se ejecuta branch_computations[N-1] como la rama predeterminada.

Cada branch_computations[b] debe tomar un solo argumento del tipo \(T_b\) y se invocará con branch_operands[b], que debe ser del mismo tipo. El tipo del valor devuelto 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.

Para obtener información sobre StableHLO, consulta StableHLO: if.

Constante

Consulta también XlaBuilder::ConstantLiteral.

Produce un output a partir de un literal constante.

Constant(literal)

Para obtener información sobre StableHLO, consulta StableHLO: constant.

ConvertElementType

Consulta también XlaBuilder::ConvertElementType.

Al igual que un static_cast por elemento en C++, ConvertElementType 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 elemento 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 float normalizadora, 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}

Para obtener información sobre StableHLO, consulta StableHLO: convert.

Conv (convolución)

Consulta también XlaBuilder::Conv.

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

Conv Pone en cola una instrucción de convolución en el cálculo, que usa los números de dimensión de convolución predeterminados sin dilatación.

El padding se especifica de forma abreviada como SAME o VALID. El padding SAME agrega ceros al tensor de entrada (lhs) para que el tensor de salida tenga la misma forma que el de entrada cuando no se tiene en cuenta el stride. El padding VALID simplemente significa que no hay padding.

Conv(lhs, rhs, window_strides, padding, feature_group_count, batch_group_count, precision_config, preferred_element_type)

Los niveles de controles disponibles para Conv son los siguientes:

• ConvWithGeneralPadding
• ConvWithGeneralDimensions
• ConvGeneral
• ConvGeneralDilated

Sea n 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, el lado derecho también es una entrada. En una red neuronal, estas son las activaciones de entrada. Las dimensiones n+2 son, en este orden:

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

El argumento rhs es un array de (n+2) dimensiones que describe el filtro, el kernel o la ventana convolucional. Las dimensiones son, en este orden:

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

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

El argumento padding especifica la cantidad de relleno con ceros que se aplicará al área base. La cantidad de padding puede ser negativa. El valor absoluto del padding 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 relleno de 2 ceros a la izquierda y de 3 ceros a la derecha en la segunda dimensión espacial. Usar relleno 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á al lado izquierdo y al lado derecho, 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 agujeros entre cada una de las entradas en esa dimensión, lo que aumenta el tamaño del array. Los agujeros se rellenan con un valor no operativo, que para la convolución significa ceros.

La dilatación del RHS también se denomina convolución atrous. Para obtener más información, consulta tf.nn.atrous_conv2d. La dilatación del LHD también se conoce como convolución transpuesta. Para obtener más información, consulta tf.nn.conv2d_transpose.

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

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

El argumento batch_group_count (valor predeterminado 1) se puede usar para los filtros agrupados durante la retropropagación. batch_group_count debe ser un divisor del tamaño de la dimensión del lote de lhs (entrada). Si batch_group_count es mayor que 1, significa que la dimensión del lote de salida debe ser de tamaño input batch / batch_group_count. El valor de batch_group_count debe ser un divisor del tamaño 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 el lado izquierdo.
• z: Tiene el mismo tamaño que output-z en el kernel (rhs).
• spatial_dims: Un valor para cada posición válida de la ventana convolucional.

En la figura anterior, se muestra cómo funciona el campo batch_group_count. De hecho, dividimos cada lote del lado izquierdo en grupos de batch_group_count y hacemos lo mismo con las características de salida. Luego, para cada uno de estos grupos, realizamos convoluciones por pares y concatenamos la salida a lo largo de la dimensión de atributos de salida. La semántica operativa de todas las demás dimensiones (espacial y de atributos) sigue siendo la misma.

Las posiciones válidas de la ventana convolucional 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 2D y elige algunas coordenadas fijas batch, z, y, x en el resultado. Luego, (y,x) es la posición de una esquina de la ventana dentro del área base (p.ej., la esquina superior izquierda, según cómo interpretes las dimensiones espaciales). Ahora tenemos una ventana 2D, tomada del área base, en la que cada punto 2D se asocia a un vector 1D, por lo que obtenemos una caja 3D. Desde el kernel convolucional, dado que fijamos la coordenada de salida z, también tenemos una caja 3D. Las dos cajas tienen las mismas dimensiones, por lo que podemos sumar los productos de los elementos correspondientes entre las dos cajas (de forma similar a un producto escalar). Ese es el valor de salida.

Ten en cuenta que, si output-z es, p.ej., 5, entonces cada posición de la ventana produce 5 valores en el resultado en la dimensión z del resultado. Estos valores difieren en qué parte del kernel convolucional se usa: hay una caja 3D separada de valores que se usa para cada coordenada output-z. Por lo tanto, podrías pensar en ella como 5 convoluciones separadas con un filtro diferente para cada una.

A continuación, se muestra el seudocódigo para 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;
}

precision_config se usa para indicar la configuración de precisión. El nivel indica si el hardware debe intentar generar más instrucciones de código de máquina para proporcionar una emulación de dtype más precisa cuando sea necesario (es decir, emular f32 en una TPU que solo admite matmuls de bf16). Los valores pueden ser DEFAULT, HIGH o HIGHEST. Encontrarás detalles adicionales en las secciones de la MXU.

preferred_element_type es un elemento escalar de tipos de salida de precisión mayor o menor que se usa para la acumulación. preferred_element_type recomienda el tipo de acumulación para la operación determinada, pero no se garantiza. Esto permite que algunos backends de hardware se acumulen en un tipo diferente y se conviertan al tipo de salida preferido.

Para obtener información sobre StableHLO, consulta StableHLO: convolución.

ConvWithGeneralPadding

Consulta también XlaBuilder::ConvWithGeneralPadding.

ConvWithGeneralPadding(lhs, rhs, window_strides, padding, feature_group_count, batch_group_count, precision_config, preferred_element_type)

Es igual que Conv, donde la configuración del padding es explícita.

ConvWithGeneralDimensions

Consulta también XlaBuilder::ConvWithGeneralDimensions.

ConvWithGeneralDimensions(lhs, rhs, window_strides, padding, dimension_numbers, feature_group_count, batch_group_count, precision_config, preferred_element_type)

Es igual que Conv, donde los números de dimensión son explícitos.

ConvGeneral

Consulta también XlaBuilder::ConvGeneral.

ConvGeneral(lhs, rhs, window_strides, padding, dimension_numbers, feature_group_count, batch_group_count, precision_config, preferred_element_type)

Es igual que Conv, donde la configuración del padding y los números de dimensión son explícitos.

ConvGeneralDilated

Consulta también XlaBuilder::ConvGeneralDilated.

ConvGeneralDilated(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, dimension_numbers, feature_group_count, batch_group_count, precision_config, preferred_element_type, window_reversal)

Es igual que Conv, donde la configuración de padding, los factores de dilatación y los números de dimensión son explícitos.

Copiar

Consulta también HloInstruction::CreateCopyStart.

Internamente, Copy se descompone en 2 instrucciones de HLO, CopyStart y CopyDone. Copy, junto con CopyStart y CopyDone, sirven como primitivas en HLO. Estas operaciones pueden aparecer en volcados de HLO, pero no están diseñadas para que los usuarios finales las construyan manualmente.

COS

Consulta tambiénXlaBuilder::Cos.

Coseno de elemento inteligente x -> cos(x).

Cos(operand)

Cos también admite el argumento opcional result_accuracy:

Cos(operand, result_accuracy)

Para obtener más información sobre result_accuracy, consulta Precisión de los resultados.

Para obtener información sobre StableHLO, consulta StableHLO: coseno.

Cosh

Consulta también XlaBuilder::Cosh.

Coseno hiperbólico x -> cosh(x) según cada elemento.

Cosh(operand)

Cosh también admite el argumento opcional result_accuracy:

Cosh(operand, result_accuracy)

Para obtener más información sobre result_accuracy, consulta Precisión de los resultados.

CustomCall

Consulta también XlaBuilder::CustomCall.

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

La documentación de CustomCall se proporciona en Detalles para desarrolladores: Llamadas personalizadas de XLA

Para obtener información sobre StableHLO, consulta StableHLO - custom_call.

Div

Consulta también XlaBuilder::Div.

Realiza la división de cada elemento del dividendo lhs y el divisor rhs.

Div(lhs, rhs)

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

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

Existe una variante alternativa con compatibilidad para la transmisión de diferentes dimensiones para Div:

Div(lhs,rhs, broadcast_dimensions)

Esta variante de la operación se debe usar para operaciones aritméticas entre arrays de diferentes rangos (como sumar una matriz a un vector).

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

Para obtener información sobre StableHLO, consulta StableHLO: divide.

Dominio

Consulta también HloInstruction::CreateDomain.

Domain puede aparecer en los volcados de HLO, pero no está diseñado para que los usuarios finales lo construyan manualmente.

Punto

Consulta también XlaBuilder::Dot.

Dot(lhs, rhs, precision_config, preferred_element_type)

La semántica exacta de esta operación depende de los rangos 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 "contraídas". Las dimensiones contraídas de lhs y rhs deben tener el mismo tamaño. En la práctica, se puede usar para realizar productos escalares entre vectores, multiplicaciones de vectores o matrices, o multiplicaciones de matrices.

precision_config se usa para indicar la configuración de precisión. El nivel indica si el hardware debe intentar generar más instrucciones de código de máquina para proporcionar una emulación de dtype más precisa cuando sea necesario (es decir, emular f32 en una TPU que solo admite matmuls de bf16). Los valores pueden ser DEFAULT, HIGH o HIGHEST. Encontrarás detalles adicionales en las secciones de la MXU.

preferred_element_type es un elemento escalar de tipos de salida de precisión mayor o menor que se usa para la acumulación. preferred_element_type recomienda el tipo de acumulación para la operación determinada, pero no se garantiza. Esto permite que algunos backends de hardware se acumulen en un tipo diferente y se conviertan al tipo de salida preferido.

Para obtener información sobre StableHLO, consulta StableHLO - dot.

DotGeneral

Consulta también XlaBuilder::DotGeneral.

DotGeneral(lhs, rhs, dimension_numbers, precision_config, preferred_element_type)

Es similar a Dot, pero permite especificar números de dimensiones de contracción y de lote para lhs y rhs.

DotGeneral realiza la suma de productos en las dimensiones de contracció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 dimensiones contraídas:

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 de lote (tamaño del lote 2, matrices de 2 x 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} } }

De ello se deduce que el número de dimensión resultante comienza con la dimensión de lote, luego la dimensión lhs no contractiva/no de lote y, finalmente, la dimensión rhs no contractiva/no de lote.

precision_config se usa para indicar la configuración de precisión. El nivel indica si el hardware debe intentar generar más instrucciones de código de máquina para proporcionar una emulación de dtype más precisa cuando sea necesario (es decir, emular f32 en una TPU que solo admite matmuls de bf16). Los valores pueden ser DEFAULT, HIGH o HIGHEST. Encontrarás más detalles en las secciones de la MXU.

preferred_element_type es un elemento escalar de tipos de salida de precisión mayor o menor que se usa para la acumulación. preferred_element_type recomienda el tipo de acumulación para la operación determinada, pero no se garantiza. Esto permite que algunos backends de hardware se acumulen en un tipo diferente y se conviertan al tipo de salida preferido.

Para obtener información sobre StableHLO, consulta StableHLO - dot_general.

ScaledDot

Consulta también XlaBuilder::ScaledDot.

ScaledDot(lhs, lhs_scale, rhs, rhs_scale, dimension_number, precision_config,preferred_element_type)

Es similar a DotGeneral.

Crea una operación de punto escalada con los operandos "lhs", "lhs_scale", "rhs" y "rhs_scale", con las dimensiones de contracción y lote especificadas en "dimension_numbers".

RaggedDot

Consulta también XlaBuilder::RaggedDot.

Para obtener un desglose del cálculo de RaggedDot, consulta StableHLO - chlo.ragged_dot.

DynamicReshape

Consulta también XlaBuilder::DynamicReshape.

Esta operación es funcionalmente idéntica a reshape, pero la forma del resultado se especifica de forma dinámica a través de output_shape.

DynamicReshape(operand, dim_sizes, new_size_bounds, dims_are_dynamic)

Para obtener información sobre StableHLO, consulta StableHLO - dynamic_reshape.

DynamicSlice

Consulta también XlaBuilder::DynamicSlice.

DynamicSlice extrae un subarreglo del arreglo de entrada en start_indices dinámico. El tamaño de la segmentación en cada dimensión se pasa en size_indices, que especifica el extremo de los intervalos de segmentación exclusivos en cada dimensión: [inicio, inicio + tamaño). La forma de start_indices debe ser unidimensional, con un tamaño de dimensión igual a la cantidad de dimensiones de operand.

DynamicSlice(operand, start_indices, slice_sizes)

Los índices de segmentación efectivos se calculan aplicando la siguiente transformación para cada índice i en [1, N) antes de realizar la segmentación:

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

Esto garantiza que el segmento extraído siempre esté dentro de los límites con respecto al array de operandos. Si el segmento está dentro de los límites antes de que se aplique la transformación, esta no tendrá ningún efecto.

Ejemplo unidimensional:

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

DynamicSlice(a, s, {2});
// Result: {2.0, 3.0}

Ejemplo bidimensional:

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

DynamicSlice(b, s, {2, 2});
//Result:
// { { 7.0,  8.0},
//   {10.0, 11.0} }

Para obtener información sobre StableHLO, consulta StableHLO - dynamic_slice.

DynamicUpdateSlice

Consulta también XlaBuilder::DynamicUpdateSlice.

DynamicUpdateSlice genera un resultado que es el valor del array de entrada operand, con un segmento update sobrescrito en start_indices. La forma de update determina la forma del subarreglo 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 segmentación efectivos se calculan aplicando la siguiente transformación para cada índice i en [1, N) antes de realizar la segmentación:

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

Esto garantiza que el segmento actualizado siempre esté dentro de los límites con respecto al array de operandos. Si el segmento está dentro de los límites antes de que se aplique la transformación, esta no tendrá ningún efecto.

Ejemplo unidimensional:

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

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

Ejemplo bidimensional:

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

let s = {1, 1}

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

Para obtener información sobre StableHLO, consulta StableHLO - dynamic_update_slice.

Erf

Consulta también XlaBuilder::Erf.

Función error x -> erf(x) para cada elemento, donde:

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

Erf(operand)

Erf también admite el argumento opcional result_accuracy:

Erf(operand, result_accuracy)

Para obtener más información sobre result_accuracy, consulta Precisión de los resultados.

Exponencial

Consulta también XlaBuilder::Exp.

Exponencial natural x -> e^x para cada elemento.

Exp(operand)

Exp también admite el argumento opcional result_accuracy:

Exp(operand, result_accuracy)

Para obtener más información sobre result_accuracy, consulta Precisión de los resultados.

Para obtener información sobre StableHLO, consulta StableHLO: exponencial.

Expm1

Consulta también XlaBuilder::Expm1.

Exponencial natural menos uno x -> e^x - 1 para cada elemento.

Expm1(operand)

La función Expm1 también admite el argumento opcional result_accuracy:

Expm1(operand, result_accuracy)

Para obtener más información sobre result_accuracy, consulta Precisión de los resultados.

Para obtener información sobre StableHLO, consulta StableHLO - exponential_minus_one.

Fft

Consulta también XlaBuilder::Fft.

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

Fft(operand, ftt_type, fft_length)

Para obtener información sobre StableHLO, consulta StableHLO - 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->complejo y complejo->real, la transformación del eje más interno se realiza (de manera efectiva) primero (FFT real a compleja; última para la FFT inversa de compleja a real), por lo que el eje más interno es el que cambia de tamaño. Las otras transformaciones de ejes serán de complejo a complejo.

Detalles de implementación

La FFT de CPU se basa en TensorFFT de Eigen. La FFT de la GPU usa cuFFT.

Piso

Consulta también XlaBuilder::Floor.

Función de piso para cada elemento x -> ⌊x⌋.

Floor(operand)

Para obtener información sobre StableHLO, consulta StableHLO: piso.

Fusión

Consulta también HloInstruction::CreateFusion.

La operación Fusion representa instrucciones de HLO y sirve como primitiva en HLO. Esta operación puede aparecer en volcados de HLO, pero no está diseñada para que los usuarios finales la construyan manualmente.

Recopila

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

Para obtener información sobre StableHLO, consulta StableHLO: gather.

Semántica general

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

gather(operand, start_indices, dimension_numbers, slice_sizes, indices_are_sorted)

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

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

Los límites para el 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, y omitimos index_vector_dim (es decir, elegimos start_indices.shape.dims[k] si k < index_vector_dim y start_indices.shape.dims[k+1] en caso 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).

Formalmente, el índice de operando In correspondiente 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 segmentar un vector S de modo que S[i] = start_indices[Combine(G, i)], donde Combine(A, b) inserta b en la posición index_vector_dim en A. Ten en cuenta que esto está bien definido incluso si G está vacío: si G está vacío, entonces S = start_indices.

2. Crea un índice inicial, Sin, en operand con S dispersando S con start_index_map. Más precisamente:

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

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

3. Crea un índice Oin en operand dispersando los índices en las dimensiones de desplazamiento en Out según el conjunto collapsed_slice_dims. Más precisamente:

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

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

4. In es Oin + Sin, donde + es la suma de cada 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 ejemplo, 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 lo son, la semántica se define según la implementación.

Descripción y ejemplos informales

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 en start_indices.

• Usamos start_index_map para correlacionar el índice inicial (cuyo tamaño puede ser menor que operand.rank) con un índice inicial "completo" en operand.

• Realizamos un corte dinámico de un segmento con un tamaño de slice_sizes usando el índice de inicio completo.

• Redimensionamos el segmento contrayendo las dimensiones collapsed_slice_dims. Dado que todas las dimensiones de segmentación contraídas deben tener un límite de 1, este cambio de forma siempre es válido.

• Usamos las dimensiones de desplazamiento en Out para indexar esta divisió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 comprender cómo se relacionan todos los elementos anteriores, veamos un ejemplo que recopila 5 segmentos de forma [8,6] de un array [16,11]. La posición de un segmento en el array [16,11] se puede representar como un vector de índice de forma S64[2], por lo que el conjunto de 5 posiciones se puede representar como un array S64[5,2].

Luego, 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 del segmento.

Esta operación de recopilación actúa como un segmento dinámico por lotes con G como dimensión de 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" con la forma [4,5,2] traduciría los índices de la siguiente manera:

Nuevamente, esto actúa como un segmento dinámico por lotes G0 y G1 como las dimensiones del lote. El tamaño de la segmentació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 desplazamiento (las dimensiones que contienen O0 y 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 desplazamiento.

2. La cantidad de dimensiones de desplazamiento 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 indican explícitamente como collapsed_slice_dims, deben tener un tamaño de segmentación de 1. Dado que tienen un tamaño de segmento de 1, el único índice válido para ellos es 0, y omitirlos no genera ambigüedad.

3. Es posible que el segmento extraído del array "Gather Indices" ((X, Y) en el último ejemplo) tenga menos elementos que la cantidad de dimensiones del array de entrada, y una asignación explícita indica 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 segmentar un índice inicial del array de índices de recopilación como de costumbre, excepto que el índice inicial 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, se expanden de acuerdo con "Asignación de índice de recopilación" (start_index_map en la descripción formal) y "Asignación de desplazamiento" (remapped_offset_dims en la descripción formal) en [X,0] y [0,O0], respectivamente, lo que suma [X,O0]. En otras palabras, el índice de salida [G0,G1,O0] se asigna al índice de entrada [GatherIndices[G0,G1,0],O0], lo que nos da la semántica para tf.gather_nd.

slice_sizes para este caso es [1,11]. Intuitivamente, 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.

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

GetDimensionSize(operand, dimension)

Para obtener información sobre StableHLO, consulta StableHLO - get_dimension_size.

GetTupleElement

Consulta también XlaBuilder::GetTupleElement.

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

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

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

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

Consulta también tf.tuple.

GetTupleElement(tuple_data, index)

Para obtener información sobre StableHLO, consulta StableHLO - get_tuple_element.

Imag

Consulta también XlaBuilder::Imag.

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

Imag(operand)

Para obtener información sobre StableHLO, consulta StableHLO: imag.

In-feed

Consulta también XlaBuilder::Infeed.

Infeed(shape, config)

Lee un solo elemento de datos de la interfaz de transmisión Infeed implícita del dispositivo, interpreta los datos como la forma y el diseño determinados, y devuelve un XlaOp de los datos. Se permiten varias operaciones de Infeed en un cálculo, pero debe haber un orden total entre ellas. Por ejemplo, dos Infeed 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 tuplas anidadas. Para una forma de tupla vacía, la operación Infeed es, en efecto, una operación nula y continúa sin leer ningún dato de la Infeed del dispositivo.

Para obtener información sobre StableHLO, consulta StableHLO: Infeed.

Pizca

Consulta también XlaBuilder::Iota.

Iota(shape, iota_dimension)

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

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

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

Para obtener información sobre StableHLO, consulta StableHLO: Iota.

IsFinite

Consulta también XlaBuilder::IsFinite.

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

IsFinite(operand)

Para obtener información sobre StableHLO, consulta StableHLO - is_finite.

Registro

Consulta también XlaBuilder::Log.

Logaritmo natural x -> ln(x) según cada elemento.

Log(operand)

El registro también admite el argumento opcional result_accuracy:

Log(operand, result_accuracy)

Para obtener más información sobre result_accuracy, consulta Precisión de los resultados.

Para obtener información sobre StableHLO, consulta StableHLO: log.

Log1p

Consulta también XlaBuilder::Log1p.

Logaritmo natural desplazado x -> ln(1+x) para cada elemento.

Log1p(operand)

Log1p también admite el argumento opcional result_accuracy:

Log1p(operand, result_accuracy)

Para obtener más información sobre result_accuracy, consulta Precisión de los resultados.

Para obtener información sobre StableHLO, consulta StableHLO: log_plus_one.

Logística

Consulta también XlaBuilder::Logistic.

Cálculo de la función logística para cada elemento x -> logistic(x).

Logistic(operand)

Logistic también admite el argumento opcional result_accuracy:

Logistic(operand, result_accuracy)

Para obtener más información sobre result_accuracy, consulta Precisión de los resultados.

Para obtener información sobre StableHLO, consulta StableHLO: Logística.

Mapa

Consulta también XlaBuilder::Map.

Map(operands..., computation, dimensions)

Aplica una función escalar a los arrays operands proporcionados, lo que genera 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. La salida tiene las mismas dimensiones que los operandos, excepto que el tipo de elemento T se reemplaza por S.

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

Para obtener información sobre StableHLO, consulta StableHLO: mapa.

Máx.

Consulta también XlaBuilder::Max.

Realiza la operación de máximo para cada elemento en los tensores lhs y rhs.

Max(lhs, rhs)

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

Existe una variante alternativa con compatibilidad para la transmisión de diferentes dimensiones para Max:

Max(lhs,rhs, broadcast_dimensions)

Esta variante de la operación se debe usar para operaciones aritméticas entre arrays de diferentes rangos (como sumar una matriz a un vector).

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

Para obtener información sobre StableHLO, consulta StableHLO: máximo.

Mín.

Consulta también XlaBuilder::Min.

Realiza una operación min en cada elemento de lhs y rhs.

Min(lhs, rhs)

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

Existe una variante alternativa con compatibilidad para la transmisión de diferentes dimensiones para Min:

Min(lhs,rhs, broadcast_dimensions)

Esta variante de la operación se debe usar para operaciones aritméticas entre arrays de diferentes rangos (como sumar una matriz a un vector).

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

Para obtener información sobre StableHLO, consulta StableHLO: mínimo.

Mul

Consulta también XlaBuilder::Mul.

Realiza el producto de los elementos de lhs y rhs.

Mul(lhs, rhs)

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

Existe una variante alternativa con compatibilidad para la transmisión de diferentes dimensiones para Mul:

Mul(lhs,rhs, broadcast_dimensions)

Esta variante de la operación se debe usar para operaciones aritméticas entre arrays de diferentes rangos (como sumar una matriz a un vector).

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

Para obtener información sobre StableHLO, consulta StableHLO: multiply.

Neg

Consulta también XlaBuilder::Neg.

Negación de cada elemento x -> -x.

Neg(operand)

Para obtener información sobre StableHLO, consulta StableHLO: negate.

No

Consulta también XlaBuilder::Not.

Es la negación lógica de cada elemento x -> !(x).

Not(operand)

Para obtener información sobre StableHLO, consulta StableHLO: No.

OptimizationBarrier

Consulta también XlaBuilder::OptimizationBarrier.

Impide que cualquier paso de optimización mueva los cálculos a través de la barrera.

OptimizationBarrier(operand)

Garantiza que todas las entradas se evalúen antes de cualquier operador que dependa de los resultados de la barrera.

Para obtener información sobre StableHLO, consulta StableHLO: optimization_barrier.

O

Consulta también XlaBuilder::Or.

Realiza una operación OR a nivel del elemento de lhs y rhs .

Or(lhs, rhs)

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

Existe una variante alternativa con compatibilidad para la transmisión de diferentes dimensiones para Or:

Or(lhs,rhs, broadcast_dimensions)

Esta variante de la operación se debe usar para operaciones aritméticas entre arrays de diferentes rangos (como sumar una matriz a un vector).

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

Para obtener información sobre StableHLO, consulta StableHLO: o.

Salida

Consulta también XlaBuilder::Outfeed.

Escribe las entradas en el feed externo.

Outfeed(operand, shape_with_layout, outfeed_config)

shape_with_layout comunica la forma diseñada que queremos alimentar.

Para obtener información sobre StableHLO, consulta StableHLO: Outfeed.

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 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 deben quitar de la dimensión especificada.

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

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

Para obtener información sobre StableHLO, consulta StableHLO: pad.

Parámetro

Consulta también XlaBuilder::Parameter.

Parameter representa un argumento de entrada para un cálculo.

PartitionID

Consulta también XlaBuilder::BuildPartitionId.

Produce partition_id del proceso actual.

PartitionID(shape)

PartitionID puede aparecer en los volcados de HLO, pero no está diseñado para que los usuarios finales lo construyan de forma manual.

Para obtener información sobre StableHLO, consulta StableHLO - partition_id.

PopulationCount

Consulta también XlaBuilder::PopulationCount.

Calcula la cantidad de bits establecidos en cada elemento de operand.

PopulationCount(operand)

Para obtener información sobre StableHLO, consulta StableHLO: popcnt.

Pow

Consulta también XlaBuilder::Pow.

Realiza la exponenciación de lhs por rhs según cada elemento.

Pow(lhs, rhs)

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

Existe una variante alternativa con compatibilidad para la transmisión de diferentes dimensiones para Pow:

Pow(lhs,rhs, broadcast_dimensions)

Esta variante de la operación se debe usar para operaciones aritméticas entre arrays de diferentes rangos (como sumar una matriz a un vector).

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

Para obtener información sobre StableHLO, consulta StableHLO: potencia.

Real

Consulta también XlaBuilder::Real.

Parte real de cada elemento de una forma compleja (o real). x -> real(x). Si el operando es de tipo de punto flotante, Real devuelve el mismo valor.

Real(operand)

Para obtener información sobre StableHLO, consulta StableHLO: real.

Recv

Consulta también XlaBuilder::Recv.

Recv, RecvWithTokens y RecvToHost son operaciones que sirven como primitivas de comunicación en HLO. Por lo general, estas operaciones aparecen en los volcados de HLO como parte de la entrada/salida de bajo nivel o la transferencia entre dispositivos, pero no están diseñadas para que los usuarios finales las construyan de forma manual.

Recv(shape, handle)

Recibe datos de la forma determinada a partir de una instrucción Send en otro cálculo que comparte el mismo identificador de canal. Devuelve un XlaOp para los datos recibidos.

Para obtener información sobre StableHLO, consulta StableHLO - recv.

RecvDone

Consulta también HloInstruction::CreateRecv y HloInstruction::CreateRecvDone.

Al igual que Send, la API de 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.

Recv(const Shape& shape, int64 channel_id)

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

Dado un contexto creado por una instrucción Recv, espera a que se complete la transferencia de datos y devuelve 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_to_reduce)

Donde:

• N debe ser mayor o igual que 1.
• El cálculo 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 N > 1, Collate(T_0, ..., T_{N-1}) es una tupla de elementos N del tipo T.

Esta operación reduce una o más dimensiones de cada array de entrada a escalares. La cantidad de dimensiones de cada array devuelto es number_of_dimensions(operand) - len(dimensions). El resultado de la operación es Collate(Q_0, ..., Q_N), donde Q_i es un array de tipo T_i, cuyas dimensiones se describen a continuación.

Se permite que diferentes backends reasocien el cálculo de la reducción. Esto puede generar diferencias numéricas, ya que algunas funciones de reducción, como la suma, no son asociativas para los números de punto flotante. Sin embargo, si el rango de los datos es limitado, la suma de números de punto flotante es lo suficientemente asociativa para la mayoría de los usos prácticos.

Para obtener información sobre StableHLO, consulta StableHLO: reduce.

Ejemplos

Cuando se reduce en una dimensión en un solo array unidimensional con valores [10, 11, 12, 13], con la función de reducción f (esto es computation), se podría calcular de la siguiente manera:

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

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

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

A continuación, se muestra un ejemplo aproximado de pseudocódigo sobre 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]

A continuación, se muestra un ejemplo de cómo reducir un array bidimensional (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 las dimensiones 0 o 1 con una función "add":

Ten en cuenta que ambos resultados de reducción son arrays unidimensionales. El diagrama muestra uno como columna y otro como fila solo para mayor comodidad visual.

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

De manera similar al ejemplo en 2D, podemos reducir solo una dimensión. Si reducimos la dimensión 0, por ejemplo, obtenemos un array bidimensional en el que todos los valores de la dimensión 0 se plegaron en un escalar:

|  4   8  12 |
| 16  20  24 |

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

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

Ten en cuenta que el orden relativo entre las dimensiones restantes en la entrada se conserva en la salida, 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. La reducción de la suma de las dimensiones 0 y 1 produce el array unidimensional [20, 28, 36].

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

Reducción variádica

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

• Valor reducido acumulado para el primer operando
• …
• Valor reducido en ejecución para el N-ésimo operando
• Valor de entrada del primer operando
• …
• Valor de entrada del N-ésimo operando

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 unidimensional 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 de inicialización I_V = Float, I_K = Int, el resultado f_(N-1) de la reducción en la única dimensión de entrada es equivalente a la siguiente aplicación recursiva:

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

Si se aplica esta reducción a un array de valores y a un array de índices secuenciales (es decir, iota), se realizará una coiteración sobre los arrays y se devolverá 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 del exponente y la 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, exponent_bits, mantissa_bits)

El resultado es un array de tipo T. Los valores de entrada se redondean al valor representable más cercano con la cantidad determinada de bits de mantisa (con la semántica "empates a par"), y los valores que superan el rango especificado por la cantidad de bits de exponente se ajustan a infinito positivo o negativo. Los valores de NaN se conservan, aunque se pueden convertir en valores de 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 del exponente o la mantisa puede exceder el valor correspondiente para el tipo T; la porción correspondiente de la conversión es, entonces, simplemente una operación nula.

Para obtener información sobre StableHLO, consulta StableHLO: reduce_precision.

ReduceScatter

Consulta también XlaBuilder::ReduceScatter.

ReduceScatter es una operación colectiva que realiza de forma eficaz un AllReduce y, luego, dispersa el resultado dividiéndolo en bloques de 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_dimension, shard_count, replica_groups, channel_id, layout, use_global_device_ids)

• Cuando operand es una tupla de arrays, la operación de reducción y 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 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 una reducción entre las réplicas 0 y 2, y 1 y 3, y, luego, dispersa el resultado.
• shard_count es el tamaño de cada grupo de réplicas. Necesitamos este campo 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 de reduce-scatter con el mismo channel_id pueden comunicarse entre sí.
• layout Consulta xla::shapes para obtener más información sobre los diseños.
• use_global_device_ids es una marca especificada por el usuario. Cuando es false(predeterminado), los números en replica_groups son ReplicaId cuando true, y replica_groups representa un ID global de (ReplicaID*partition_count + partition_id). Por ejemplo:
  • Con 2 réplicas y 4 particiones,
  • replica_groups={ {0,1,4,5},{2,3,6,7} } and use_global_device_ids=true
  • group[0] = (0,0), (0,1), (1,0), (1,1)
  • group[1] = (0,2), (0,3), (1,2), (1,3)
  • donde cada par es (replica_id, partition_id).

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 el valor [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.

Para obtener información sobre StableHLO, consulta StableHLO - reduce_scatter.

ReduceScatter, ejemplo 1, StableHLO

En el ejemplo anterior, hay 2 réplicas que participan en ReduceScatter. En cada réplica, el operando tiene la forma f32[2,4]. Se realiza una reducción total (suma) en todas las réplicas, lo que produce un valor reducido de la forma f32[2,4] en cada réplica. Luego, este valor reducido se divide en 2 partes a lo largo de la dimensión 1, de modo que cada parte tenga la forma f32[2,2]. Cada réplica dentro del grupo de procesos recibe la parte correspondiente a su posición en el grupo. Como resultado, la salida en cada réplica tiene la forma f32[2,2].

ReduceWindow

Consulta también XlaBuilder::ReduceWindow.

Aplica una función de reducción a todos los elementos de cada ventana de una secuencia de arrays multidimensionales de N, lo que produce un solo array multidimensional o una tupla de N arrays multidimensionales como salida. Cada array de salida tiene la misma cantidad de elementos que la cantidad de posiciones válidas de la ventana. Una capa de agrupación se puede expresar como un ReduceWindow. Al igual que con Reduce, el computation aplicado siempre se pasa al init_values del 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 N > 1, Collate(T_0, ..., T_{N-1}) es una tupla de elementos N del tipo (T0,...T{N-1}).

Para obtener información sobre StableHLO, consulta StableHLO - reduce_window.

ReduceWindow, ejemplo 1

La entrada es una matriz de tamaño [4x6], y tanto window_dimensions como window_stride_dimensions son [2x3].

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

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