Semântica de operação

A seguir, descrevemos a semântica das operações definidas na interface XlaBuilder. Normalmente, essas operações são mapeadas de um para um para as operações definidas na interface RPC em xla_data.proto.

Uma observação sobre a nomenclatura: o tipo de dados generalizado com que a XLA trabalha é uma matriz N-dimensional que contém elementos de algum tipo uniforme (como ponto flutuante de 32 bits). Em toda a documentação, array é usado para denotar uma matriz de dimensão arbitrária. Para facilitar, os casos especiais têm nomes mais específicos e conhecidos. Por exemplo, um vetor é uma matriz unidimensional, e uma matriz é uma matriz bidimensional.

Saiba mais sobre a estrutura de uma operação em Formas e layout e Layout em blocos.

Abdômen

Consulte também XlaBuilder::Abs.

Valor absoluto de elemento x -> |x|.

Abs(operand)

Para informações sobre o StableHLO, consulte StableHLO: abs.

Adicionar

Consulte também XlaBuilder::Add.

Executa a adição de elementos de lhs e rhs.

Add(lhs, rhs)

As formas dos argumentos precisam ser semelhantes ou compatíveis. Consulte a documentação de transmissão para saber o que significa a compatibilidade de formas. O resultado de uma operação tem uma forma que é o resultado da transmissão das duas matrizes de entrada. Nessa variante, as operações entre matrizes de ranks diferentes não são aceitas, a menos que um dos operandos seja um escalar.

Existe uma variante alternativa com suporte de transmissão de diferentes dimensões para "Add":

Add(lhs,rhs, broadcast_dimensions)

Essa variante da operação deve ser usada para operações aritméticas entre matrizes de ranks diferentes (como adicionar uma matriz a um vetor).

O operando broadcast_dimensions adicional é uma fração de números inteiros que especifica as dimensões a serem usadas para transmissão dos operandos. A semântica é descrita em detalhes na página de transmissão.

Para informações sobre o StableHLO, consulte StableHLO: adicionar.

AddDependency

Consulte também HloInstruction::AddDependency.

AddDependency pode aparecer em despejos de HLO, mas não foi criado para ser construído manualmente por usuários finais.

AfterAll

Consulte também XlaBuilder::AfterAll.

"AfterAll" usa um número variádico de tokens e produz um único token. Os tokens são tipos primitivos que podem ser encadeados entre operações com efeitos colaterais para impor a ordenação. AfterAll pode ser usado como uma junção de tokens para ordenar uma operação após um conjunto de operações.

AfterAll(tokens)

Para informações sobre o StableHLO, consulte StableHLO - after_all.

AllGather

Consulte também XlaBuilder::AllGather.

Realiza a concatenação entre réplicas.

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

• replica_groups é uma lista de grupos de réplicas entre os quais a concatenação é realizada. O ID da réplica atual pode ser recuperado usando ReplicaId. A ordem das réplicas em cada grupo determina a ordem em que as entradas são localizadas no resultado. replica_groups precisa estar vazio (nesse caso, todas as réplicas pertencem a um único grupo, ordenado de 0 a N - 1) ou conter o mesmo número de elementos que o número de réplicas. Por exemplo, replica_groups = {0, 2}, {1, 3} realiza a concatenação entre as réplicas 0 e 2, e 1 e 3.
• shard_count é o tamanho de cada grupo de réplicas. Precisamos disso nos casos em que replica_groups estão vazios.
• channel_id é usado para comunicação entre módulos: apenas operações all-gather com o mesmo channel_id podem se comunicar entre si.
• use_global_device_ids Retorna "true" se os IDs na configuração do ReplicaGroup representarem um ID global de (replica_id * partition_count + partition_id) em vez de um ID de réplica. Isso permite um agrupamento mais flexível de dispositivos se essa redução total for entre partições e entre réplicas.

O formato de saída é o formato de entrada com o all_gather_dimension shard_count vezes maior. Por exemplo, se houver duas réplicas e o operando tiver o valor [1.0, 2.5] e [3.0, 5.25], respectivamente, nas duas réplicas, o valor de saída dessa operação em que all_gather_dim é 0 será [1.0, 2.5, 3.0,5.25] nas duas réplicas.

A API de AllGather é decomposta internamente em duas instruções HLO (AllGatherStart e AllGatherDone).

Consulte também HloInstruction::CreateAllGatherStart.

AllGatherStart e AllGatherDone servem como primitivos em HLO. Essas operações podem aparecer em despejos de HLO, mas não devem ser criadas manualmente por usuários finais.

Para informações sobre o StableHLO, consulte StableHLO: all_gather.

AllReduce

Consulte também XlaBuilder::AllReduce.

Realiza uma computação personalizada em várias réplicas.

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

• Quando operand é uma tupla de matrizes, a redução completa é realizada em cada elemento da tupla.
• replica_groups é uma lista de grupos de réplicas entre os quais a redução é realizada. O ID da réplica atual pode ser recuperado usando ReplicaId. replica_groups precisa estar vazio (nesse caso, todas as réplicas pertencem a um único grupo) ou conter o mesmo número de elementos que o número de réplicas. Por exemplo, replica_groups = {0, 2}, {1, 3} realiza a redução entre as réplicas 0 e 2, e 1 e 3.
• channel_id é usado para comunicação entre módulos: apenas operações all-reduce com o mesmo channel_id podem se comunicar entre si.
• shape_with_layout: força o layout do AllReduce para o layout especificado. Isso é usado para garantir o mesmo layout para um grupo de operações AllReduce compiladas separadamente.
• use_global_device_ids Retorna "true" se os IDs na configuração do ReplicaGroup representarem um ID global de (replica_id * partition_count + partition_id) em vez de um ID de réplica. Isso permite um agrupamento mais flexível de dispositivos se essa redução total for entre partições e entre réplicas.

O formato da saída é igual ao da entrada. Por exemplo, se houver duas réplicas e o operando tiver o valor [1.0, 2.5] e [3.0, 5.25], respectivamente, nas duas réplicas, o valor de saída dessa operação e o cálculo de soma serão [4.0, 7.75] nas duas réplicas. Se a entrada for uma tupla, a saída também será uma tupla.

Para calcular o resultado de AllReduce, é necessário ter uma entrada de cada réplica. Portanto, se uma réplica executar um nó AllReduce mais vezes do que outra, a primeira vai esperar para sempre. Como todas as réplicas executam o mesmo programa, não há muitas maneiras de isso acontecer. No entanto, é possível quando a condição de um loop "while" depende de dados de infeed e os dados que são infeed fazem com que o loop "while" itere mais vezes em uma réplica do que em outra.

A API de AllReduce é decomposta internamente em duas instruções HLO (AllReduceStart e AllReduceDone).

Consulte também HloInstruction::CreateAllReduceStart.

AllReduceStart e AllReduceDone servem como primitivos em HLO. Essas operações podem aparecer em despejos de HLO, mas não devem ser criadas manualmente por usuários finais.

CrossReplicaSum

Consulte também XlaBuilder::CrossReplicaSum.

Realiza AllReduce com um cálculo de soma.

CrossReplicaSum(operand, replica_groups)

Retorna a soma do valor do operando em cada subgrupo de réplicas. Todas as réplicas fornecem uma entrada para a soma e recebem o resultado para cada subgrupo.

AllToAll

Consulte também XlaBuilder::AllToAll.

AllToAll é uma operação coletiva que envia dados de todos os núcleos para todos os núcleos. Ele tem duas fases:

1. A fase de dispersão. Em cada núcleo, o operando é dividido em split_count blocos ao longo da split_dimensions, e os blocos são distribuídos para todos os núcleos. Por exemplo, o iésimo bloco é enviado para o iésimo núcleo.
2. A fase de coleta. Cada núcleo concatena os blocos recebidos ao longo da concat_dimension.

Os núcleos participantes podem ser configurados por:

• replica_groups: cada ReplicaGroup contém uma lista de IDs de réplicas que participam da computação. O ID da réplica atual pode ser recuperado usando ReplicaId. O AllToAll será aplicado em subgrupos na ordem especificada. Por exemplo, replica_groups = { {1,2,3}, {4,5,0} } significa que um AllToAll será aplicado em réplicas {1, 2, 3}, na fase de coleta, e os blocos recebidos serão concatenados na mesma ordem de 1, 2, 3. Em seguida, outro AllToAll será aplicado nas réplicas 4, 5 e 0, e a ordem de concatenação também será 4, 5 e 0. Se replica_groups estiver vazio, todas as réplicas vão pertencer a um grupo, na ordem de concatenação da aparência delas.

Pré-requisitos:

• O tamanho da dimensão do operando em split_dimension é divisível por split_count.
• O formato do operando não é uma tupla.

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

Consulte xla::shapes para mais informações sobre formas e layouts.

Para informações sobre o StableHLO, consulte StableHLO: all_to_all.

AllToAll: exemplo 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);

No exemplo acima, há quatro núcleos participando do Alltoall. Em cada núcleo, o operando é dividido em quatro partes ao longo da dimensão 1, de modo que cada parte tem formato f32[4,4]. As quatro partes são distribuídas para todos os núcleos. Em seguida, cada núcleo concatena as partes recebidas ao longo da dimensão 0, na ordem dos núcleos 0 a 4. Portanto, a saída em cada núcleo tem o formato f32[16,4].

AllToAll - Exemplo 2 - StableHLO

No exemplo acima, há duas réplicas participando do AllToAll. Em cada réplica, o operando tem a forma f32[2,4]. O operando é dividido em duas partes ao longo da dimensão 1, então cada parte tem a forma f32[2,2]. As duas partes são trocadas entre as réplicas de acordo com a posição delas no grupo. Cada réplica coleta a parte correspondente dos dois operandos e os concatena ao longo da dimensão 0. Como resultado, a saída em cada réplica tem o formato f32[4,2].

RaggedAllToAll

Consulte também XlaBuilder::RaggedAllToAll.

O RaggedAllToAll realiza uma operação coletiva de todos para todos, em que a entrada e a saída são tensores irregulares.

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

Tensores irregulares são definidos por um conjunto de três tensores:

• data: o tensor data é irregular ao longo da dimensão mais externa, em que cada elemento indexado tem tamanho variável.
• offsets: o tensor offsets indexa a dimensão mais externa do tensor data e representa o deslocamento inicial de cada elemento irregular do tensor data.
• sizes: o tensor sizes representa o tamanho de cada elemento irregular do tensor data, em que o tamanho é especificado em unidades de subelementos. Um subelemento é definido como o sufixo da forma do tensor "data" obtido removendo a dimensão mais externa.
• Os tensores offsets e sizes precisam ter o mesmo tamanho.

Exemplo de um 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 precisa ser fragmentado de forma que cada réplica tenha offsets na perspectiva de saída da réplica de destino.

Para o i-ésimo deslocamento de saída, a réplica atual vai enviar a atualização input[input_offsets[i]:input_offsets[i]+input_sizes[i]] para a i-ésima réplica, que será gravada em output_i[output_offsets[i]:output_offsets[i]+send_sizes[i]] na i-ésima réplica output.

Por exemplo, se tivermos duas 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]

O HLO irregular de todos para todos tem os seguintes argumentos:

• input: tensor de dados de entrada irregulares.
• output: tensor de dados de saída irregulares.
• input_offsets: tensor de deslocamentos de entrada irregulares.
• send_sizes: tensor irregular de tamanhos de envio.
• output_offsets: matriz de intervalos irregulares na saída da réplica de destino.
• recv_sizes: tensor de tamanhos de recebimento irregulares.

Os tensores *_offsets e *_sizes precisam ter o mesmo formato.

Há duas formas compatíveis para os tensores *_offsets e *_sizes:

• [num_devices] em que ragged-all-to-all pode enviar no máximo uma atualização para cada dispositivo remoto no grupo de réplicas. Exemplo:

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] em que ragged-all-to-all pode enviar até num_updates atualizações para o mesmo dispositivo remoto (cada uma em diferentes offsets), para cada dispositivo remoto no grupo de réplicas.

Exemplo:

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

E

Consulte também XlaBuilder::And.

Executa um AND bit a bit de dois tensores lhs e rhs.

And(lhs, rhs)

As formas dos argumentos precisam ser semelhantes ou compatíveis. Consulte a documentação de transmissão para saber o que significa a compatibilidade de formas. O resultado de uma operação tem uma forma que é o resultado da transmissão das duas matrizes de entrada. Nessa variante, as operações entre matrizes de ranks diferentes não são aceitas, a menos que um dos operandos seja um escalar.

Existe uma variante alternativa com suporte de transmissão de dimensões diferentes para "And":

And(lhs,rhs, broadcast_dimensions)

Essa variante da operação deve ser usada para operações aritméticas entre matrizes de ranks diferentes (como adicionar uma matriz a um vetor).

O operando broadcast_dimensions adicional é uma fração de números inteiros que especifica as dimensões a serem usadas para transmissão dos operandos. A semântica é descrita em detalhes na página de transmissão.

Para informações sobre o StableHLO, consulte StableHLO - and.

Assíncrona

Consulte também HloInstruction::CreateAsyncStart, HloInstruction::CreateAsyncUpdate, HloInstruction::CreateAsyncDone.

AsyncDone, AsyncStart e AsyncUpdate são instruções HLO internas usadas para operações assíncronas e servem como primitivos em HLO. Essas operações podem aparecer em despejos de HLO, mas não devem ser construídas manualmente por usuários finais.

Atan2

Consulte também XlaBuilder::Atan2.

Executa a operação atan2 elemento a elemento em lhs e rhs.

Atan2(lhs, rhs)

As formas dos argumentos precisam ser semelhantes ou compatíveis. Consulte a documentação de transmissão para saber o que significa a compatibilidade de formas. O resultado de uma operação tem uma forma que é o resultado da transmissão das duas matrizes de entrada. Nessa variante, as operações entre matrizes de ranks diferentes não são aceitas, a menos que um dos operandos seja um escalar.

Existe uma variante alternativa com suporte de transmissão de dimensões diferentes para Atan2:

Atan2(lhs,rhs, broadcast_dimensions)

Essa variante da operação deve ser usada para operações aritméticas entre matrizes de ranks diferentes (como adicionar uma matriz a um vetor).

O operando broadcast_dimensions adicional é uma fração de números inteiros que especifica as dimensões a serem usadas para transmissão dos operandos. A semântica é descrita em detalhes na página de transmissão.

Para informações sobre o StableHLO, consulte StableHLO: atan2.

BatchNormGrad

Consulte também XlaBuilder::BatchNormGrad e o artigo original sobre normalização em lote para uma descrição detalhada do algoritmo.

Calcula gradientes de normalização em lote.

BatchNormGrad(operand, scale, batch_mean, batch_var, grad_output, epsilon, feature_index)

Para cada recurso na dimensão de recurso (feature_index é o índice da dimensão de recurso em operand), a operação calcula os gradientes em relação a operand, offset e scale em todas as outras dimensões. O feature_index precisa ser um índice válido para a dimensão de recurso em operand.

Os três gradientes são definidos pelas seguintes fórmulas (supondo uma matriz de quatro dimensões como operand e com índice de dimensão de recurso l, tamanho do lote m e tamanhos espaciais w e 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} \]

As entradas batch_mean e batch_var representam valores de momentos em dimensões espaciais e de lote.

O tipo de saída é uma tupla de três identificadores:

Para informações sobre o StableHLO, consulte StableHLO - batch_norm_grad.

BatchNormInference

Consulte também XlaBuilder::BatchNormInference e o artigo original sobre normalização em lote para uma descrição detalhada do algoritmo.

Normaliza uma matriz em lote e dimensões espaciais.

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

Para cada recurso na dimensão de recurso (feature_index é o índice da dimensão de recurso em operand), a operação calcula a média e a variância em todas as outras dimensões e usa esses valores para normalizar cada elemento em operand. O feature_index precisa ser um índice válido para a dimensão de recurso em operand.

BatchNormInference é equivalente a chamar BatchNormTraining sem calcular mean e variance para cada lote. Em vez disso, ele usa as entradas mean e variance como valores estimados. O objetivo dessa operação é reduzir a latência na inferência, daí o nome BatchNormInference.

A saída é uma matriz normalizada n-dimensional com o mesmo formato da entrada operand.

Para informações sobre o StableHLO, consulte StableHLO - batch_norm_inference.

BatchNormTraining

Consulte também XlaBuilder::BatchNormTraining e the original batch normalization paper para uma descrição detalhada do algoritmo.

Normaliza uma matriz em lote e dimensões espaciais.

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

Para cada recurso na dimensão de recurso (feature_index é o índice da dimensão de recurso em operand), a operação calcula a média e a variância em todas as outras dimensões e usa esses valores para normalizar cada elemento em operand. O feature_index precisa ser um índice válido para a dimensão de recurso em operand.

O algoritmo funciona da seguinte maneira para cada lote em operand \(x\) que contém melementos com w e h como o tamanho das dimensões espaciais (supondo que operandseja uma matriz de quatro dimensões):

• Calcula a média do lote \(\mu_l\) para cada recurso l na dimensão do recurso: \(\mu_l=\frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h x_{ijkl}\)

• Calcula a variância do 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, dimensiona e muda: \(y_{ijkl}=\frac{\gamma_l(x_{ijkl}-\mu_l)}{\sqrt[2]{\sigma^2_l+\epsilon} }+\beta_l\)

O valor de epsilon, geralmente um número pequeno, é adicionado para evitar erros de divisão por zero.

O tipo de saída é uma tupla de três XlaOps:

Os batch_mean e batch_var são momentos calculados em todas as dimensões espaciais e de lote usando as fórmulas acima.

Para informações sobre o StableHLO, consulte StableHLO - batch_norm_training.

Bitcast

Consulte também HloInstruction::CreateBitcast.

Bitcast pode aparecer em despejos de HLO, mas não foi criado para ser construído manualmente por usuários finais.

BitcastConvertType

Consulte também XlaBuilder::BitcastConvertType.

Semelhante a um tf.bitcast no TensorFlow, realiza uma operação de bitcast de elemento a elemento de uma forma de dados para uma forma de destino. O tamanho da entrada e da saída precisa ser igual. Por exemplo, elementos s32 se tornam elementos f32 usando a rotina bitcast, e um elemento s32 se torna quatro elementos s8. O Bitcast é implementado como uma transmissão de baixo nível. Portanto, máquinas com representações de ponto flutuante diferentes vão gerar resultados diferentes.

BitcastConvertType(operand, new_element_type)

As dimensões do operando e da forma de destino precisam corresponder, exceto a última dimensão, que vai mudar pela proporção do tamanho primitivo antes e depois da conversão.

Os tipos de elementos de origem e destino não podem ser tuplas.

Para mais informações sobre o StableHLO, consulte StableHLO - bitcast_convert.

Conversão de bitcast para tipo primitivo de largura diferente

A instrução BitcastConvert HLO é compatível com o caso em que o tamanho do tipo de elemento de saída T' não é igual ao tamanho do elemento de entrada T. Como toda a operação é conceitualmente um bitcast e não muda os bytes subjacentes, a forma do elemento de saída precisa mudar. Para B = sizeof(T), B' = sizeof(T'), há dois casos possíveis.

Primeiro, quando B > B', o formato da saída recebe uma nova dimensão secundária de tamanho B/B'. Exemplo:

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

A regra permanece a mesma para escalares efetivos:

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

Como alternativa, para B' > B, a instrução exige que a última dimensão lógica da forma de entrada seja igual a B'/B, e essa dimensão é descartada durante a conversão:

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

As conversões entre diferentes larguras de bits não são feitas elemento a elemento.

Transmitir

Consulte também XlaBuilder::Broadcast.

Adiciona dimensões a uma matriz duplicando os dados nela.

Broadcast(operand, broadcast_sizes)

As novas dimensões são inseridas à esquerda. Por exemplo, se broadcast_sizes tiver valores {a0, ..., aN} e a forma do operando tiver dimensões {b0, ..., bM}, a forma da saída terá dimensões {a0, ..., aN, b0, ..., bM}.

O novo índice de dimensões em cópias do operando, ou seja,

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

Por exemplo, se operand for um escalar f32 com valor 2.0f, e broadcast_sizes for {2, 3}, o resultado será uma matriz com formato f32[2, 3] e todos os valores no resultado serão 2.0f.

Para informações sobre o StableHLO, consulte StableHLO: transmissão.

BroadcastInDim

Consulte também XlaBuilder::BroadcastInDim.

Aumenta o tamanho e o número de dimensões de uma matriz duplicando os dados nela.

BroadcastInDim(operand, out_dim_size, broadcast_dimensions)

Semelhante à transmissão, mas permite adicionar dimensões em qualquer lugar e expandir as dimensões atuais com tamanho 1.

O operand é transmitido para a forma descrita por out_dim_size. broadcast_dimensions mapeia as dimensões de operand para as dimensões da forma de destino. Ou seja, a i-ésima dimensão do operando é mapeada para a dimensão broadcast_dimension[i] da forma de saída. As dimensões de operand precisam ter tamanho 1 ou o mesmo tamanho da dimensão na forma de saída a que são mapeadas. As dimensões restantes são preenchidas com dimensões de tamanho 1. A transmissão de dimensões degeneradas é feita ao longo dessas dimensões para alcançar a forma de saída. A semântica é descrita em detalhes na página de transmissão.

Ligar

Consulte também XlaBuilder::Call.

Invoca um cálculo com os argumentos fornecidos.

Call(computation, operands...)

A aridade e os tipos de operands precisam corresponder aos parâmetros de computation. É permitido não ter operands.

CompositeCall

Consulte também XlaBuilder::CompositeCall.

Encapsula uma operação composta por outras operações do StableHLO, recebendo entradas e composite_attributes e produzindo resultados. A semântica da operação é implementada pelo atributo de decomposição. A operação composta pode ser substituída pela decomposição sem mudar a semântica do programa. Nos casos em que a incorporação da decomposição não fornece a mesma semântica de operação, prefira usar custom_call.

O campo de versão (padrão é 0) é usado para indicar quando a semântica de um composto muda.

Essa operação é implementada como um kCall com o atributo is_composite=true. O campo decomposition é especificado pelo atributo computation. Os atributos do front-end armazenam os atributos restantes prefixados com composite..

Exemplo de operação 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)

O decomposition de uma operação não é um campo chamado, mas aparece como um atributo to_apply que aponta para a função que contém a implementação de nível inferior, ou seja, to_apply=%funcname.

Para mais informações sobre composição e decomposição, consulte a especificação do StableHLO.

Cbrt

Consulte também XlaBuilder::Cbrt.

Operação de raiz cúbica com elementos x -> cbrt(x).

Cbrt(operand)

O Cbrt também é compatível com o argumento opcional result_accuracy:

Cbrt(operand, result_accuracy)

Para mais informações sobre result_accuracy, consulte Precisão dos resultados.

Para informações sobre o StableHLO, consulte StableHLO: cbrt.

Arredondar para cima

Consulte também XlaBuilder::Ceil.

Teto de elemento a elemento x -> ⌈x⌉.

Ceil(operand)

Para informações sobre o StableHLO, consulte StableHLO: ceil.

Cholesky

Consulte também XlaBuilder::Cholesky.

Calcula a decomposição de Cholesky de um lote de matrizes simétricas (hermitianas) positivas definidas.

Cholesky(a, lower)

Se lower for true, vai calcular matrizes triangulares inferiores l de forma que $a = l . l^T$. Se lower for false, vai calcular matrizes triangulares superiores u de forma que\(a = u^T . u\).

Os dados de entrada são lidos apenas do triângulo inferior/superior de a, dependendo do valor de lower. Os valores do outro triângulo são ignorados. Os dados de saída são retornados no mesmo triângulo. Os valores no outro triângulo são definidos pela implementação e podem ser qualquer coisa.

Se a tiver mais de duas dimensões, ele será tratado como um lote de matrizes, em que todas, exceto as duas dimensões secundárias, são dimensões de lote.a

Se a não for simétrica (hermitiana) positiva definida, o resultado será definido pela implementação.

Para informações sobre o StableHLO, consulte StableHLO: cholesky.

Limitar

Consulte também XlaBuilder::Clamp.

Restringe um operando ao intervalo entre um valor mínimo e máximo.

Clamp(min, operand, max)

Dado um operando e valores mínimo e máximo, retorna o operando se ele estiver no intervalo entre o mínimo e o máximo. Caso contrário, retorna o valor mínimo se o operando estiver abaixo desse intervalo ou o valor máximo se o operando estiver acima desse intervalo. Ou seja, clamp(a, x, b) = min(max(a, x), b).

Todos os três precisam ter o mesmo formato. Como alternativa, como uma forma restrita de transmissão, min e/ou max podem ser um escalar do tipo T.

Exemplo com min e 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 informações sobre o StableHLO, consulte StableHLO: clamp.

Recolher

Consulte também XlaBuilder::Collapse. e a operação tf.reshape.

Reduz as dimensões de uma matriz em uma dimensão.

Collapse(operand, dimensions)

A redução substitui o subconjunto especificado das dimensões do operando por uma única dimensão. Os argumentos de entrada são uma matriz arbitrária do tipo T e um vetor constante de tempo de compilação de índices de dimensão. Os índices de dimensão precisam ser um subconjunto ordenado (números de dimensão de baixo para cima) e consecutivo das dimensões de T. Portanto, {0, 1, 2}, {0, 1} ou {1, 2} são conjuntos de dimensões válidos, mas {1, 0} ou {0, 2} não são. Elas são substituídas por uma única dimensão nova, na mesma posição na sequência de dimensões que as substituídas, com o tamanho da nova dimensão igual ao produto dos tamanhos das dimensões originais. O menor número de dimensão em dimensions é a dimensão de variação mais lenta (mais importante) no aninhamento de loop que recolhe essas dimensões, e o maior número de dimensão é a variação mais rápida (mais secundária). Consulte o operador tf.reshape se for necessário um ordenamento de redução mais geral.

Por exemplo, seja v uma matriz 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

Consulte também XlaBuilder::Clz.

Contagem de zeros à esquerda por elemento.

Clz(operand)

CollectiveBroadcast

Consulte também XlaBuilder::CollectiveBroadcast.

Transmite dados entre réplicas. Os dados são enviados do primeiro ID de réplica em cada grupo para os outros IDs no mesmo grupo. Se um ID de réplica não estiver em nenhum grupo de réplicas, a saída nessa réplica será um tensor composto por 0s em shape.

CollectiveBroadcast(operand, replica_groups, channel_id)

Para informações sobre o StableHLO, consulte StableHLO: collective_broadcast.

CollectivePermute

Consulte também XlaBuilder::CollectivePermute.

CollectivePermute é uma operação coletiva que envia e recebe dados entre réplicas.

CollectivePermute(operand, source_target_pairs, channel_id, inplace)

Há as seguintes restrições no source_target_pairs:

• Dois pares não podem ter o mesmo ID de réplica de destino nem de origem.
• Se um ID de réplica não for um destino em nenhum par, a saída nessa réplica será um tensor composto de 0s com o mesmo formato da entrada.

A API da operação CollectivePermute é decomposta internamente em duas instruções HLO (CollectivePermuteStart e CollectivePermuteDone).

Consulte também HloInstruction::CreateCollectivePermuteStart.

CollectivePermuteStart e CollectivePermuteDone servem como primitivos em HLO. Essas operações podem aparecer em despejos de HLO, mas não devem ser construídas manualmente por usuários finais.

Para informações sobre o StableHLO, consulte StableHLO - collective_permute.

Comparar

Consulte também XlaBuilder::Compare.

Executa uma comparação elemento a elemento de lhs e rhs do seguinte:

Eq

Consulte também XlaBuilder::Eq.

Realiza uma comparação igual a por elemento de lhs e rhs.

\(lhs = rhs\)

Eq(lhs, rhs)

As formas dos argumentos precisam ser semelhantes ou compatíveis. Consulte a documentação de transmissão para saber o que significa a compatibilidade de formas. O resultado de uma operação tem uma forma que é o resultado da transmissão das duas matrizes de entrada. Nessa variante, as operações entre matrizes de ranks diferentes não são aceitas, a menos que um dos operandos seja um escalar.

Existe uma variante alternativa com suporte de transmissão de dimensões diferentes para a equação:

Eq(lhs,rhs, broadcast_dimensions)

Essa variante da operação deve ser usada para operações aritméticas entre matrizes de ranks diferentes (como adicionar uma matriz a um vetor).

O operando broadcast_dimensions adicional é uma fração de números inteiros que especifica as dimensões a serem usadas para transmissão dos operandos. A semântica é descrita em detalhes na página de transmissão.

Existe suporte para uma ordem total sobre os números de ponto flutuante para Eq, aplicando:

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

EqTotalOrder(lhs,rhs, broadcast_dimensions)

Para informações sobre o StableHLO, consulte StableHLO: comparação.

Ne

Consulte também XlaBuilder::Ne.

Executa uma comparação diferente de de elemento a elemento entre lhs e rhs.

\(lhs != rhs\)

Ne(lhs, rhs)

As formas dos argumentos precisam ser semelhantes ou compatíveis. Consulte a documentação de transmissão para saber o que significa a compatibilidade de formas. O resultado de uma operação tem uma forma que é o resultado da transmissão das duas matrizes de entrada. Nessa variante, as operações entre matrizes de ranks diferentes não são aceitas, a menos que um dos operandos seja um escalar.

Existe uma variante alternativa com suporte de transmissão de dimensões diferentes para Ne:

Ne(lhs,rhs, broadcast_dimensions)

Essa variante da operação deve ser usada para operações aritméticas entre matrizes de ranks diferentes (como adicionar uma matriz a um vetor).

O operando broadcast_dimensions adicional é uma fração de números inteiros que especifica as dimensões a serem usadas para transmissão dos operandos. A semântica é descrita em detalhes na página de transmissão.

Existe suporte para uma ordem total sobre os números de ponto flutuante para Ne, aplicando:

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

NeTotalOrder(lhs,rhs, broadcast_dimensions)

Para informações sobre o StableHLO, consulte StableHLO: comparação.

Ge

Consulte também XlaBuilder::Ge.

Executa uma comparação greater-or-equal-than elemento a elemento de lhs e rhs.

\(lhs >= rhs\)

Ge(lhs, rhs)

As formas dos argumentos precisam ser semelhantes ou compatíveis. Consulte a documentação de transmissão para saber o que significa a compatibilidade de formas. O resultado de uma operação tem uma forma que é o resultado da transmissão das duas matrizes de entrada. Nessa variante, as operações entre matrizes de ranks diferentes não são aceitas, a menos que um dos operandos seja um escalar.

Existe uma variante alternativa com suporte de transmissão de dimensões diferentes para Ge:

Ge(lhs,rhs, broadcast_dimensions)

Essa variante da operação deve ser usada para operações aritméticas entre matrizes de ranks diferentes (como adicionar uma matriz a um vetor).

O operando broadcast_dimensions adicional é uma fração de números inteiros que especifica as dimensões a serem usadas para transmissão dos operandos. A semântica é descrita em detalhes na página de transmissão.

Existe uma ordem total sobre os números de ponto flutuante para Gt, aplicando:

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

GtTotalOrder(lhs,rhs, broadcast_dimensions)

Para informações sobre o StableHLO, consulte StableHLO: comparação.

Gt

Consulte também XlaBuilder::Gt.

Executa uma comparação maior que de elemento a elemento de lhs e rhs.

\(lhs > rhs\)

Gt(lhs, rhs)

As formas dos argumentos precisam ser semelhantes ou compatíveis. Consulte a documentação de transmissão para saber o que significa a compatibilidade de formas. O resultado de uma operação tem uma forma que é o resultado da transmissão das duas matrizes de entrada. Nessa variante, as operações entre matrizes de ranks diferentes não são aceitas, a menos que um dos operandos seja um escalar.

Existe uma variante alternativa com suporte de transmissão de diferentes dimensões para Gt:

Gt(lhs,rhs, broadcast_dimensions)

Essa variante da operação deve ser usada para operações aritméticas entre matrizes de ranks diferentes (como adicionar uma matriz a um vetor).

O operando broadcast_dimensions adicional é uma fração de números inteiros que especifica as dimensões a serem usadas para transmissão dos operandos. A semântica é descrita em detalhes na página de transmissão.

Para informações sobre o StableHLO, consulte StableHLO: comparação.

Le

Consulte também XlaBuilder::Le.

Executa uma comparação less-or-equal-than de elemento a elemento de lhs e rhs.

\(lhs <= rhs\)

Le(lhs, rhs)

As formas dos argumentos precisam ser semelhantes ou compatíveis. Consulte a documentação de transmissão para saber o que significa a compatibilidade de formas. O resultado de uma operação tem uma forma que é o resultado da transmissão das duas matrizes de entrada. Nessa variante, as operações entre matrizes de ranks diferentes não são aceitas, a menos que um dos operandos seja um escalar.

Existe uma variante alternativa com suporte de transmissão de dimensões diferentes para Le:

Le(lhs,rhs, broadcast_dimensions)

Essa variante da operação deve ser usada para operações aritméticas entre matrizes de ranks diferentes (como adicionar uma matriz a um vetor).

O operando broadcast_dimensions adicional é uma fração de números inteiros que especifica as dimensões a serem usadas para transmissão dos operandos. A semântica é descrita em detalhes na página de transmissão.

Existe uma ordem total sobre os números de ponto flutuante para Le, aplicando:

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

LeTotalOrder(lhs,rhs, broadcast_dimensions)

Para informações sobre o StableHLO, consulte StableHLO: comparação.

Lt

Consulte também XlaBuilder::Lt.

Executa uma comparação menor que elemento a elemento de lhs e rhs.

\(lhs < rhs\)

Lt(lhs, rhs)

As formas dos argumentos precisam ser semelhantes ou compatíveis. Consulte a documentação de transmissão para saber o que significa a compatibilidade de formas. O resultado de uma operação tem uma forma que é o resultado da transmissão das duas matrizes de entrada. Nessa variante, as operações entre matrizes de ranks diferentes não são aceitas, a menos que um dos operandos seja um escalar.

Existe uma variante alternativa com suporte de transmissão de dimensões diferentes para Lt:

Lt(lhs,rhs, broadcast_dimensions)

Essa variante da operação deve ser usada para operações aritméticas entre matrizes de ranks diferentes (como adicionar uma matriz a um vetor).

O operando broadcast_dimensions adicional é uma fração de números inteiros que especifica as dimensões a serem usadas para transmissão dos operandos. A semântica é descrita em detalhes na página de transmissão.

Suporte a um pedido total sobre os números de ponto flutuante para Lt, aplicando:

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

LtTotalOrder(lhs,rhs, broadcast_dimensions)

Para informações sobre o StableHLO, consulte StableHLO: comparação.

Complexo

Consulte também XlaBuilder::Complex.

Realiza a conversão de elemento por elemento para um valor complexo de um par de valores reais e imaginários, lhs e rhs.

Complex(lhs, rhs)

As formas dos argumentos precisam ser semelhantes ou compatíveis. Consulte a documentação de transmissão para saber o que significa a compatibilidade de formas. O resultado de uma operação tem uma forma que é o resultado da transmissão das duas matrizes de entrada. Nessa variante, as operações entre matrizes de ranks diferentes não são aceitas, a menos que um dos operandos seja um escalar.

Existe uma variante alternativa com suporte de transmissão de dimensões diferentes para números complexos:

Complex(lhs,rhs, broadcast_dimensions)

Essa variante da operação deve ser usada para operações aritméticas entre matrizes de ranks diferentes (como adicionar uma matriz a um vetor).

O operando broadcast_dimensions adicional é uma fração de números inteiros que especifica as dimensões a serem usadas para transmissão dos operandos. A semântica é descrita em detalhes na página de transmissão.

Para informações sobre o StableHLO, consulte StableHLO: complex.

ConcatInDim (concatenação)

Consulte também XlaBuilder::ConcatInDim.

A função CONCATENAR cria uma matriz com vários operandos de matriz. A matriz tem o mesmo número de dimensões que cada um dos operandos de matriz de entrada (que precisam ter o mesmo número de dimensões entre si) e contém os argumentos na ordem em que foram especificados.

Concatenate(operands..., dimension)

Com exceção de dimension, todas as dimensões precisam ser iguais. Isso acontece porque o XLA não é compatível com matrizes irregulares. Além disso, valores de dimensão zero não podem ser concatenados, já que é impossível nomear a dimensão ao longo da qual a concatenação ocorre.

Exemplo unidimensional:

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

Exemplo 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 mais informações sobre o StableHLO, consulte StableHLO: concatenar.

Condicional

Consulte também XlaBuilder::Conditional.

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

Executa true_computation se predicate for true, false_computation se predicate for false e retorna o resultado.

O true_computation precisa receber um único argumento do tipo \(T_0\) e será invocado com true_operand, que precisa ser do mesmo tipo. O false_computation precisa usar um único argumento do tipo \(T_1\) e será invocado com false_operand, que precisa ser do mesmo tipo. O tipo do valor retornado de true_computation e false_computation precisa ser o mesmo.

Apenas um de true_computation e false_computation será executado, dependendo do valor de predicate.

Conditional(branch_index, branch_computations, branch_operands)

Executa branch_computations[branch_index] e retorna o resultado. Se branch_index for um S32 que seja < 0 ou >= N, branch_computations[N-1] será executado como a ramificação padrão.

Cada branch_computations[b] precisa receber um único argumento do tipo \(T_b\) e será invocado com branch_operands[b], que precisa ser do mesmo tipo. O tipo do valor retornado de cada branch_computations[b] precisa ser o mesmo.

Apenas um dos branch_computations será executado, dependendo do valor de branch_index.

Para informações sobre o StableHLO, consulte StableHLO: if.

Constante

Consulte também XlaBuilder::ConstantLiteral.

Produz um output de um literal constante.

Constant(literal)

Para informações sobre o StableHLO, consulte StableHLO: constante.

ConvertElementType

Consulte também XlaBuilder::ConvertElementType.

Semelhante a um static_cast independente em C++, o ConvertElementType realiza uma operação de conversão independente de uma forma de dados para uma forma de destino. As dimensões precisam corresponder, e a conversão é elemento a elemento. Por exemplo, elementos s32 se tornam elementos f32 por uma rotina de conversão de s32 para f32.

ConvertElementType(operand, new_element_type)

As dimensões do operando e do formato de destino precisam ser iguais. Os tipos de elementos de origem e destino não podem ser tuplas.

Uma conversão como T=s32 para U=f32 vai executar uma rotina de conversão de int para float normalizadora, como arredondamento para o número par mais próximo.

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 informações sobre o StableHLO, consulte StableHLO: converter.

Conv (convolução)

Consulte também XlaBuilder::Conv.

Calcula uma convolução do tipo usado em redes neurais. Aqui, uma convolução pode ser considerada como uma janela n-dimensional que se move por uma área de base n-dimensional, e um cálculo é realizado para cada posição possível da janela.

Conv Enfileira uma instrução de convolução no cálculo, que usa os números de dimensão de convolução padrão sem dilatação.

O padding é especificado de forma abreviada como SAME ou VALID. O padding SAME preenche a entrada (lhs) com zeros para que a saída tenha o mesmo formato da entrada quando não considera o stride. O preenchimento VALID significa que não há preenchimento.

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

Há níveis crescentes de controles disponíveis para Conv:

• ConvWithGeneralPadding
• ConvWithGeneralDimensions
• ConvGeneral
• ConvGeneralDilated

Seja n o número de dimensões espaciais. O argumento lhs é uma matriz (n+2)-dimensional que descreve a área de base. Isso é chamado de entrada, embora o lado direito também seja uma entrada. Em uma rede neural, essas são as ativações de entrada. As dimensões n+2 são, nesta ordem:

• batch: cada coordenada nessa dimensão representa uma entrada independente para a qual a convolução é realizada.
• z/depth/features: cada posição (y,x) na área de base tem um vetor associado a ela, que entra nessa dimensão.
• spatial_dims: descreve as dimensões espaciais n que definem a área base em que a janela se move.

O argumento rhs é uma matriz (n+2)-dimensional que descreve o filtro/kernel/janela convolucional. As dimensões são, nesta ordem:

• output-z: a dimensão z da saída.
• input-z: o tamanho dessa dimensão vezes feature_group_count precisa ser igual ao tamanho da dimensão z no lado esquerdo.
• spatial_dims: descreve as dimensões espaciais n que definem a janela n-d que se move pela área de base.

O argumento window_strides especifica a extensão da janela convolucional nas dimensões espaciais. Por exemplo, se o stride na primeira dimensão espacial for 3, a janela só poderá ser colocada em coordenadas em que o primeiro índice espacial seja divisível por 3.

O argumento padding especifica a quantidade de padding de zero a ser aplicada à área de base. A quantidade de padding pode ser negativa. O valor absoluto do padding negativo indica o número de elementos a serem removidos da dimensão especificada antes de fazer a convolução. padding[0] especifica o padding para a dimensão y, e padding[1] especifica o padding para a dimensão x. Cada par tem o padding baixo como o primeiro elemento e o padding alto como o segundo. O padding baixo é aplicado na direção de índices menores, enquanto o padding alto é aplicado na direção de índices maiores. Por exemplo, se padding[1] for (2,3), haverá um padding de dois zeros à esquerda e três zeros à direita na segunda dimensão espacial. Usar padding é equivalente a inserir esses mesmos valores zero na entrada (lhs) antes de fazer a convolução.

Os argumentos lhs_dilation e rhs_dilation especificam o fator de dilatação a ser aplicado ao lado esquerdo e direito, respectivamente, em cada dimensão espacial. Se o fator de dilatação em uma dimensão espacial for d, d-1 furos serão implicitamente colocados entre cada uma das entradas nessa dimensão, aumentando o tamanho da matriz. Os buracos são preenchidos com um valor nulo, que para convolução significa zeros.

A dilatação do rhs também é chamada de convolução atrous. Para mais detalhes, consulte tf.nn.atrous_conv2d. A dilatação do lado esquerdo também é chamada de convolução transposta. Para mais detalhes, consulte tf.nn.conv2d_transpose.

O argumento feature_group_count (valor padrão 1) pode ser usado para convoluções agrupadas. feature_group_count precisa ser um divisor da dimensão de recurso de entrada e de saída. Se feature_group_count for maior que 1, isso significa que, conceitualmente, a dimensão do recurso de entrada e saída e a dimensão do recurso de saída rhs serão divididas igualmente em muitos grupos feature_group_count, cada um consistindo em uma subsequência consecutiva de recursos. A dimensão do recurso de entrada de rhs precisa ser igual à dimensão do recurso de entrada de lhs dividida por feature_group_count. Assim, ela já tem o tamanho de um grupo de recursos de entrada. Os grupos i-ésimos são usados juntos para calcular feature_group_count para muitas convoluções separadas. Os resultados dessas convoluções são concatenados na dimensão da característica de saída.

Para a convolução por profundidade, o argumento feature_group_count seria definido como a dimensão do recurso de entrada, e o filtro seria redimensionado de [filter_height, filter_width, in_channels, channel_multiplier] para [filter_height, filter_width, 1, in_channels * channel_multiplier]. Para mais detalhes, consulte tf.nn.depthwise_conv2d.

O argumento batch_group_count (valor padrão 1) pode ser usado para filtros agrupados durante a backpropagation. batch_group_count precisa ser um divisor do tamanho da dimensão do lote lhs (entrada). Se batch_group_count for maior que 1, isso significa que a dimensão do lote de saída deve ter o tamanho input batch / batch_group_count. O batch_group_count precisa ser um divisor do tamanho do recurso de saída.

O formato da saída tem estas dimensões, nesta ordem:

• batch: o tamanho dessa dimensão vezes batch_group_count precisa ser igual ao tamanho da dimensão batch em lhs.
• z: mesmo tamanho que output-z no kernel (rhs).
• spatial_dims: um valor para cada posição válida da janela convolucional.

A figura acima mostra como o campo batch_group_count funciona. Efetivamente, dividimos cada lote do lado esquerdo em batch_group_count grupos e fazemos o mesmo com os recursos de saída. Em seguida, para cada um desses grupos, fazemos convoluções aos pares e concatenamos a saída ao longo da dimensão do recurso de saída. A semântica operacional de todas as outras dimensões (recurso e espacial) permanece a mesma.

As posições válidas da janela convolucional são determinadas pelas strides e pelo tamanho da área da base após o padding.

Para descrever o que uma convolução faz, considere uma convolução 2D e escolha algumas coordenadas fixas batch, z, y, x na saída. Então, (y,x) é uma posição de um canto da janela dentro da área de base (por exemplo, o canto superior esquerdo, dependendo de como você interpreta as dimensões espaciais). Agora temos uma janela 2D, extraída da área de base, em que cada ponto 2D está associado a um vetor 1D, resultando em uma caixa 3D. Do kernel convolucional, como fixamos a coordenada de saída z, também temos uma caixa 3D. As duas caixas têm as mesmas dimensões. Portanto, podemos somar os produtos elemento a elemento entre elas (semelhante a um produto escalar). Esse é o valor de saída.

Se output-z for, por exemplo, 5, cada posição da janela produz cinco valores na saída para a dimensão z. Esses valores diferem na parte do kernel convolucional usada. Há uma caixa 3D separada de valores usada para cada coordenada output-z. Então, pense nisso como cinco convoluções separadas com um filtro diferente para cada uma delas.

Confira um pseudocódigo para uma convolução 2D com padding e 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 é usado para indicar a configuração de precisão. O nível determina se o hardware deve tentar gerar mais instruções de código de máquina para fornecer uma emulação de dtype mais precisa quando necessário (ou seja, emular f32 em uma TPU que só aceita bf16 matmuls). Os valores podem ser DEFAULT, HIGH, HIGHEST. Mais detalhes nas seções da MXU.

preferred_element_type é um elemento escalar de tipos de saída de precisão maior/menor usado para acumulação. preferred_element_type recomenda o tipo de acumulação para a operação especificada, mas não há garantia. Isso permite que alguns backends de hardware acumulem em um tipo diferente e convertam para o tipo de saída preferido.

Para informações sobre o StableHLO, consulte StableHLO: convolução.

ConvWithGeneralPadding

Consulte também XlaBuilder::ConvWithGeneralPadding.

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

Igual a Conv, em que a configuração de padding é explícita.

ConvWithGeneralDimensions

Consulte também XlaBuilder::ConvWithGeneralDimensions.

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

Igual a Conv, em que os números de dimensão são explícitos.

ConvGeneral

Consulte também XlaBuilder::ConvGeneral.

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

Igual a Conv, em que os números de dimensão e a configuração de padding são explícitos.

ConvGeneralDilated

Consulte também 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)

Igual a Conv, em que a configuração de padding, os fatores de dilatação e os números de dimensão são explícitos.

Copiar

Consulte também HloInstruction::CreateCopyStart.

O Copy é decomposto internamente em duas instruções HLO, CopyStart e CopyDone. Copy, CopyStart e CopyDone servem como primitivos em HLO. Essas operações podem aparecer em despejos de HLO, mas não devem ser construídas manualmente por usuários finais.

Cos

Consulte também XlaBuilder::Cos.

Cosseno elemento a elemento x -> cos(x).

Cos(operand)

A função Cos também aceita o argumento opcional result_accuracy:

Cos(operand, result_accuracy)

Para mais informações sobre result_accuracy, consulte Precisão dos resultados.

Para informações sobre o StableHLO, consulte StableHLO: cosseno.

Cosh

Consulte também XlaBuilder::Cosh.

Cosseno hiperbólico x -> cosh(x) por elemento.

Cosh(operand)

A função Cosh também é compatível com o argumento opcional result_accuracy:

Cosh(operand, result_accuracy)

Para mais informações sobre result_accuracy, consulte Precisão dos resultados.

CustomCall

Consulte também XlaBuilder::CustomCall.

Chamar uma função fornecida pelo usuário em um cálculo.

A documentação do CustomCall está disponível em Detalhes do desenvolvedor: chamadas personalizadas do XLA

Para informações sobre o StableHLO, consulte StableHLO: custom_call.

Div

Consulte também XlaBuilder::Div.

Realiza a divisão elemento a elemento do dividendo lhs e do divisor rhs.

Div(lhs, rhs)

O estouro de divisão de números inteiros (divisão/resto com ou sem sinal por zero ou divisão/resto com sinal de INT_SMIN com -1) produz um valor definido pela implementação.

As formas dos argumentos precisam ser semelhantes ou compatíveis. Consulte a documentação de transmissão para saber o que significa a compatibilidade de formas. O resultado de uma operação tem uma forma que é o resultado da transmissão das duas matrizes de entrada. Nessa variante, as operações entre matrizes de ranks diferentes não são aceitas, a menos que um dos operandos seja um escalar.

Existe uma variante alternativa com suporte a transmissão de diferentes dimensões para Div:

Div(lhs,rhs, broadcast_dimensions)

Essa variante da operação deve ser usada para operações aritméticas entre matrizes de ranks diferentes (como adicionar uma matriz a um vetor).

O operando broadcast_dimensions adicional é uma fração de números inteiros que especifica as dimensões a serem usadas para transmissão dos operandos. A semântica é descrita em detalhes na página de transmissão.

Para informações sobre o StableHLO, consulte StableHLO: divisão.

Domínio

Consulte também HloInstruction::CreateDomain.

Domain pode aparecer em despejos de HLO, mas não foi criado para ser construído manualmente por usuários finais.

Escalar

Consulte também XlaBuilder::Dot.

Dot(lhs, rhs, precision_config, preferred_element_type)

A semântica exata dessa operação depende das classificações dos operandos:

A operação realiza a soma dos produtos na segunda dimensão de lhs (ou a primeira, se tiver uma dimensão) e na primeira dimensão de rhs. Essas são as dimensões "contraídas". As dimensões contraídas de lhs e rhs precisam ser do mesmo tamanho. Na prática, ele pode ser usado para realizar produtos escalares entre vetores, multiplicações de vetor/matriz ou multiplicações de matriz/matriz.

precision_config é usado para indicar a configuração de precisão. O nível determina se o hardware deve tentar gerar mais instruções de código de máquina para fornecer uma emulação de dtype mais precisa quando necessário (ou seja, emular f32 em uma TPU que só aceita bf16 matmuls). Os valores podem ser DEFAULT, HIGH, HIGHEST. Mais detalhes nas seções da MXU.

preferred_element_type é um elemento escalar de tipos de saída de precisão maior/menor usado para acumulação. preferred_element_type recomenda o tipo de acumulação para a operação especificada, mas não há garantia. Isso permite que alguns backends de hardware acumulem em um tipo diferente e convertam para o tipo de saída preferido.

Para informações sobre o StableHLO, consulte StableHLO - dot.

DotGeneral

Consulte também XlaBuilder::DotGeneral.

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

Semelhante a "Dot", mas permite que números de dimensão de contração e lote sejam especificados para lhs e rhs.

O DotGeneral realiza a soma dos produtos nas dimensões de contração especificadas em dimension_numbers.

Os números de dimensão de contração associados de lhs e rhs não precisam ser iguais, mas precisam ter os mesmos tamanhos de dimensão.

Exemplo com números de dimensão de contração:

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

Os números de dimensão de lote associados de lhs e rhs precisam ter os mesmos tamanhos de dimensão.

Exemplo com números de dimensão de lote (tamanho do lote 2, matrizes 2x2):

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

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

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

DotGeneral(lhs, rhs, dnums) -> {
    { {1.0, 2.0},
      {3.0, 4.0} },
    { {5.0, 6.0},
      {7.0, 8.0} } }

Portanto, o número da dimensão resultante começa com a dimensão de lote, depois a dimensão lhs não contraída/não em lote e, por fim, a dimensão rhs não contraída/não em lote.

precision_config é usado para indicar a configuração de precisão. O nível determina se o hardware deve tentar gerar mais instruções de código de máquina para fornecer uma emulação de dtype mais precisa quando necessário (ou seja, emular f32 em uma TPU que só aceita bf16 matmuls). Os valores podem ser DEFAULT, HIGH, HIGHEST. Mais detalhes podem ser encontrados nas seções de MXU.

preferred_element_type é um elemento escalar de tipos de saída de precisão maior/menor usado para acumulação. preferred_element_type recomenda o tipo de acumulação para a operação especificada, mas não há garantia. Isso permite que alguns backends de hardware acumulem em um tipo diferente e convertam para o tipo de saída preferido.

Para informações sobre o StableHLO, consulte StableHLO: dot_general.

ScaledDot

Consulte também XlaBuilder::ScaledDot.

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

Semelhante a DotGeneral.

Cria uma operação de ponto escalonado com os operandos "lhs", "lhs_scale", "rhs" e "rhs_scale", com dimensões de contração e lote especificadas em "dimension_numbers".

RaggedDot

Consulte também XlaBuilder::RaggedDot.

Para uma análise detalhada do cálculo de RaggedDot, consulte StableHLO - chlo.ragged_dot

DynamicReshape

Consulte também XlaBuilder::DynamicReshape.

Essa operação é funcionalmente idêntica a reshape, mas o formato do resultado é especificado dinamicamente via output_shape.

DynamicReshape(operand, dim_sizes, new_size_bounds, dims_are_dynamic)

Para informações sobre o StableHLO, consulte StableHLO: dynamic_reshape.

DynamicSlice

Consulte também XlaBuilder::DynamicSlice.

O DynamicSlice extrai uma submatriz da matriz de entrada em start_indices dinâmico. O tamanho da fatia em cada dimensão é transmitido em size_indices, que especifica o ponto final dos intervalos de fatias exclusivas em cada dimensão: [início, início + tamanho). O formato de start_indices precisa ser unidimensional, com tamanho igual ao número de dimensões de operand.

DynamicSlice(operand, start_indices, slice_sizes)

Os índices de corte efetivos são calculados aplicando a seguinte transformação para cada índice i em [1, N) antes de realizar o corte:

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

Isso garante que a fração extraída esteja sempre dentro dos limites em relação à matriz de operandos. Se a fatia estiver dentro dos limites antes da transformação, ela não terá efeito.

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

Exemplo 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 informações sobre o StableHLO, consulte StableHLO: dynamic_slice.

DynamicUpdateSlice

Consulte também XlaBuilder::DynamicUpdateSlice.

O DynamicUpdateSlice gera um resultado que é o valor da matriz de entrada operand, com uma fração update substituída em start_indices. A forma de update determina a forma da submatriz do resultado que é atualizada. O formato de start_indices precisa ser unidimensional, com tamanho igual ao número de dimensões de operand.

DynamicUpdateSlice(operand, update, start_indices)

Os índices de corte efetivos são calculados aplicando a seguinte transformação para cada índice i em [1, N) antes de realizar o corte:

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

Isso garante que a fatia atualizada esteja sempre dentro dos limites em relação à matriz de operandos. Se a fatia estiver dentro dos limites antes da transformação, ela não terá efeito.

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

Exemplo 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 informações sobre o StableHLO, consulte StableHLO: dynamic_update_slice.

Erf

Consulte também XlaBuilder::Erf.

Função de erro elemento a elemento x -> erf(x) em que:

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

Erf(operand)

A função erf também é compatível com o argumento opcional result_accuracy:

Erf(operand, result_accuracy)

Para mais informações sobre result_accuracy, consulte Precisão dos resultados.

Exp

Consulte também XlaBuilder::Exp.

Exponencial natural x -> e^x por elemento.

Exp(operand)

A função exponencial também é compatível com o argumento opcional result_accuracy:

Exp(operand, result_accuracy)

Para mais informações sobre result_accuracy, consulte Precisão dos resultados.

Para informações sobre o StableHLO, consulte StableHLO: exponencial.

Expm1

Consulte também XlaBuilder::Expm1.

Exponencial natural elemento a elemento menos um x -> e^x - 1.

Expm1(operand)

A função Expm1 também aceita o argumento opcional result_accuracy:

Expm1(operand, result_accuracy)

Para mais informações sobre result_accuracy, consulte Precisão dos resultados.

Para informações sobre o StableHLO, consulte StableHLO: exponential_minus_one.

Fft

Consulte também XlaBuilder::Fft.

A operação FFT do XLA implementa as transformações de Fourier direta e inversa para entradas/saídas reais e complexas. FFTs multidimensionais em até três eixos são compatíveis.

Fft(operand, ftt_type, fft_length)

Para informações sobre o StableHLO, consulte StableHLO: fft.

FFT multidimensional

Quando mais de um fft_length é fornecido, isso equivale a aplicar uma cascata de operações de FFT a cada um dos eixos mais internos. Nos casos real->complexo e complexo->real, a transformação do eixo mais interno é (efetivamente) realizada primeiro (RFFT; última para IRFFT). Por isso, o eixo mais interno é o que muda de tamanho. Outras transformações de eixo serão complexas -> complexas.

Detalhes da implementação

A FFT da CPU é compatível com o TensorFFT do Eigen. A FFT da GPU usa cuFFT.

Andar

Consulte também XlaBuilder::Floor.

Função de piso x -> ⌊x⌋ por elemento.

Floor(operand)

Para informações sobre o StableHLO, consulte StableHLO: floor.

Fusão

Consulte também HloInstruction::CreateFusion.

A operação Fusion representa instruções HLO e serve como uma primitiva em HLO. Essa operação pode aparecer em despejos de HLO, mas não deve ser construída manualmente por usuários finais.

Coletar

A operação de coleta do XLA une várias partes (cada uma em um deslocamento de tempo de execução potencialmente diferente) de uma matriz de entrada.

Para informações sobre o StableHLO, consulte StableHLO: gather.

Semântica geral

Consulte também XlaBuilder::Gather. Para uma descrição mais intuitiva, consulte a seção "Descrição informal" abaixo.

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

Por conveniência, rotulamos as dimensões na matriz de saída que não estão em offset_dims como batch_dims.

A saída é uma matriz com batch_dims.size + offset_dims.size dimensões.

O operand.rank precisa ser igual à soma de offset_dims.size e collapsed_slice_dims.size. Além disso, slice_sizes.size precisa ser igual a operand.rank.

Se index_vector_dim for igual a start_indices.rank, vamos considerar implicitamente que start_indices tem uma dimensão 1 final. Por exemplo, se start_indices tiver o formato [6,7] e index_vector_dim for 2, vamos considerar implicitamente que o formato de start_indices é [6,7,1].

Os limites da matriz de saída ao longo da dimensão i são calculados da seguinte forma:

1. Se i estiver presente em batch_dims (ou seja, for igual a batch_dims[k] para algum k), vamos escolher os limites de dimensão correspondentes em start_indices.shape, ignorando index_vector_dim (ou seja, escolher start_indices.shape.dims[k] se k < index_vector_dim e start_indices.shape.dims[k+1] caso contrário).

2. Se i estiver presente em offset_dims (ou seja, igual a offset_dims[k] para algum k), vamos escolher o limite correspondente de slice_sizes depois de considerar collapsed_slice_dims (ou seja, vamos escolher adjusted_slice_sizes[k], em que adjusted_slice_sizes é slice_sizes com os limites nos índices collapsed_slice_dims removidos).

Formalmente, o índice do operando In correspondente a um determinado índice de saída Out é calculado da seguinte maneira:

1. Seja G = { Out[k] para k em batch_dims }. Use G para extrair um vetor S de forma que S[i] = start_indices[Combine(G, i)], em que Combine(A, b) insere b na posição index_vector_dim em A. Isso é bem definido mesmo que G esteja vazio: se G estiver vazio, S = start_indices.

2. Crie um índice inicial, Sin, em operand usando S ao dispersar S usando start_index_map. Mais precisamente:

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

   2. Sin[_] = 0 caso contrário.

3. Crie um índice Oin em operand dispersando os índices nas dimensões de deslocamento em Out de acordo com o conjunto collapsed_slice_dims. Mais precisamente:

   1. Oin[remapped_offset_dims(k)] = Out[offset_dims[k]] se k < offset_dims.size (remapped_offset_dims é definido abaixo).

   2. Oin[_] = 0 caso contrário.

4. In é Oin + Sin, em que + é a adição de elemento a elemento.

remapped_offset_dims é uma função monotônica com domínio [0, offset_dims.size) e intervalo [0, operand.rank) \ collapsed_slice_dims. Então, se, por exemplo, offset_dims.size é 4, operand.rank é 6 e collapsed_slice_dims é {0, 2}, então remapped_offset_dims é {0→1, 1→3, 2→4, 3→5}.

Se indices_are_sorted for definido como verdadeiro, o XLA poderá presumir que start_indices estão classificados (em ordem crescente, após dispersar os valores de acordo com start_index_map) pelo usuário. Caso contrário, a semântica será definida pela implementação.

Descrição informal e exemplos

De maneira informal, cada índice Out na matriz de saída corresponde a um elemento E na matriz de operandos, calculado da seguinte maneira:

• Usamos as dimensões do lote em Out para pesquisar um índice inicial de start_indices.

• Usamos start_index_map para mapear o índice inicial (cujo tamanho pode ser menor que operand.rank) para um índice inicial "completo" no operand.

• Fazemos uma segmentação dinâmica com tamanho slice_sizes usando o índice inicial completo.

• Para isso, recolhemos as dimensões collapsed_slice_dims. Como todas as dimensões de fatia recolhidas precisam ter um limite de 1, essa mudança de formato é sempre válida.

• Usamos as dimensões de deslocamento em Out para indexar essa fatia e receber o elemento de entrada, E, correspondente ao índice de saída Out.

index_vector_dim é definido como start_indices.rank - 1 em todos os exemplos a seguir. Valores mais interessantes para index_vector_dim não mudam a operação fundamentalmente, mas tornam a representação visual mais complicada.

Para entender como tudo isso se encaixa, vamos analisar um exemplo que reúne cinco fatias de forma [8,6] de uma matriz [16,11]. A posição de uma fatia na matriz [16,11] pode ser representada como um vetor de índice de forma S64[2]. Assim, o conjunto de cinco posições pode ser representado como uma matriz S64[5,2].

O comportamento da operação de coleta pode ser descrito como uma transformação de índice que usa [G,O0,O1], um índice no formato de saída, e o mapeia para um elemento na matriz de entrada da seguinte maneira:

Primeiro, selecionamos um vetor (X,Y) da matriz de índices de coleta usando G. O elemento na matriz de saída no índice [G,O0,O1] é o elemento na matriz de entrada no índice [X+O0,Y+O1].

slice_sizes é [8,6], que decide o intervalo de O₀ e O₁, e isso, por sua vez, decide os limites da fração.

Essa operação de coleta funciona como uma fração dinâmica em lote com G como a dimensão do lote.

Os índices de coleta podem ser multidimensionais. Por exemplo, uma versão mais geral do exemplo acima usando uma matriz "reunir índices" de forma [4,5,2] traduziria os índices assim:

Novamente, isso funciona como uma segmentação dinâmica em lote G0 e G1 como as dimensões do lote. O tamanho da fração ainda é [8,6].

A operação de coleta no XLA generaliza a semântica informal descrita acima das seguintes maneiras:

1. Podemos configurar quais dimensões na forma de saída são as dimensões de ajuste (dimensões que contêm O0, O1 no último exemplo). As dimensões de lote de saída (dimensões que contêm G0, G1 no último exemplo) são definidas como as dimensões de saída que não são de compensação.

2. O número de dimensões de deslocamento de saída explicitamente presentes na forma de saída pode ser menor do que o número de dimensões de entrada. Essas dimensões "ausentes", que são listadas explicitamente como collapsed_slice_dims, precisam ter um tamanho de fração de 1. Como eles têm um tamanho de fração de 1, o único índice válido é 0, e a omissão não causa ambiguidade.

3. A fatia extraída da matriz "Gather Indices" ((X, Y) no último exemplo) pode ter menos elementos do que o número de dimensões da matriz de entrada, e um mapeamento explícito determina como o índice deve ser expandido para ter o mesmo número de dimensões da entrada.

Como exemplo final, usamos (2) e (3) para implementar tf.gather_nd:

G0 e G1 são usados para extrair um índice inicial da matriz de índices de coleta como de costume, exceto que o índice inicial tem apenas um elemento, X. Da mesma forma, há apenas um índice de ajuste de saída com o valor O0. No entanto, antes de serem usados como índices na matriz de entrada, eles são expandidos de acordo com "Coletar mapeamento de índice" (start_index_map na descrição formal) e "Mapeamento de deslocamento" (remapped_offset_dims na descrição formal) em [X,0] e [0,O0], respectivamente, somando [X,O0]. Em outras palavras, o índice de saída [G0,G1,O0] é mapeado para o índice de entrada [GatherIndices[G0,G1,0],O0], o que nos dá a semântica para tf.gather_nd.

slice_sizes para este caso é [1,11]. Isso significa que cada índice X na matriz de índices de coleta escolhe uma linha inteira, e o resultado é a concatenação de todas essas linhas.

GetDimensionSize

Consulte também XlaBuilder::GetDimensionSize.

Retorna o tamanho da dimensão especificada do operando. O operando precisa ter o formato de matriz.

GetDimensionSize(operand, dimension)

Para informações sobre o StableHLO, consulte StableHLO - get_dimension_size.

GetTupleElement

Consulte também XlaBuilder::GetTupleElement.

Indexa uma tupla com um valor constante de tempo de compilação.

O valor precisa ser uma constante de tempo de compilação para que a inferência de forma possa determinar o tipo do valor resultante.

Isso é análogo a std::get<int N>(t) em C++. Conceitualmente:

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.

Consulte também tf.tuple.

GetTupleElement(tuple_data, index)

Para informações sobre o StableHLO, consulte StableHLO: get_tuple_element.

Imag

Consulte também XlaBuilder::Imag.

Parte imaginária de um elemento de uma forma complexa (ou real). x -> imag(x). Se o operando for um tipo de ponto flutuante, será retornado 0.

Imag(operand)

Para informações sobre o StableHLO, consulte StableHLO - imag.

No feed

Consulte também XlaBuilder::Infeed.

Infeed(shape, config)

Lê um único item de dados da interface de streaming infeed implícita do dispositivo, interpretando os dados como a forma e o layout especificados, e retorna um XlaOp dos dados. Várias operações de feed são permitidas em um cálculo, mas é necessário haver uma ordem total entre elas. Por exemplo, dois Infeeds no código abaixo têm uma ordem total porque há uma dependência entre os loops "while".

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

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

Não é possível usar tuplas aninhadas. Para uma forma de tupla vazia, a operação Infeed é efetivamente uma operação nula e continua sem ler dados do Infeed do dispositivo.

Para informações sobre o StableHLO, consulte StableHLO - infeed.

Iota

Consulte também XlaBuilder::Iota.

Iota(shape, iota_dimension)

Cria uma constante literal no dispositivo em vez de uma transferência de host potencialmente grande. Cria uma matriz com o formato especificado e contém valores que começam em zero e aumentam em um ao longo da dimensão especificada. Para tipos de ponto flutuante, a matriz produzida é equivalente a ConvertElementType(Iota(...)), em que Iota é do tipo integral e a conversão é para o tipo de ponto flutuante.

Por exemplo, Iota(s32[4, 8], 0) retorna

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

Devolução 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 informações sobre o StableHLO, consulte StableHLO: iota.

IsFinite

Consulte também XlaBuilder::IsFinite.

Testa se cada elemento de operand é finito, ou seja, não é infinito positivo ou negativo e não é NaN. Retorna uma matriz de valores PRED com o mesmo formato da entrada, em que cada elemento é true se e somente se o elemento de entrada correspondente for finito.

IsFinite(operand)

Para informações sobre o StableHLO, consulte StableHLO: is_finite.

Registro

Consulte também XlaBuilder::Log.

Logaritmo natural elemento a elemento x -> ln(x).

Log(operand)

O registro também é compatível com o argumento opcional result_accuracy:

Log(operand, result_accuracy)

Para mais informações sobre result_accuracy, consulte Precisão dos resultados.

Para mais informações sobre o StableHLO, consulte StableHLO: log.

Log1p

Consulte também XlaBuilder::Log1p.

Logaritmo natural x -> ln(1+x) com deslocamento elemento a elemento.

Log1p(operand)

O Log1p também aceita o argumento opcional result_accuracy:

Log1p(operand, result_accuracy)

Para mais informações sobre result_accuracy, consulte Precisão dos resultados.

Para informações sobre o StableHLO, consulte StableHLO - log_plus_one.

Logística

Consulte também XlaBuilder::Logistic.

Cálculo da função logística elemento a elemento x -> logistic(x).

Logistic(operand)

A logística também aceita o argumento opcional result_accuracy:

Logistic(operand, result_accuracy)

Para mais informações sobre result_accuracy, consulte Precisão dos resultados.

Para informações sobre o StableHLO, consulte StableHLO: logística.

Mapa

Consulte também XlaBuilder::Map.

Map(operands..., computation, dimensions)

Aplica uma função escalar às matrizes operands fornecidas, produzindo uma matriz de mesmas dimensões em que cada elemento é o resultado da função mapeada aplicada aos elementos correspondentes nas matrizes de entrada.

A função mapeada é um cálculo arbitrário com a restrição de ter N entradas do tipo escalar T e uma única saída do tipo S. A saída tem as mesmas dimensões que os operandos, exceto que o tipo de elemento T é substituído por S.

Por exemplo, Map(op1, op2, op3, computation, par1) mapeia elem_out <- computation(elem1, elem2, elem3, par1) em cada índice (multidimensional) nas matrizes de entrada para produzir a matriz de saída.

Para informações sobre o StableHLO, consulte StableHLO: mapa.

Máx.

Consulte também XlaBuilder::Max.

Executa a operação de máximo por elemento nos tensores lhs e rhs.

Max(lhs, rhs)

As formas dos argumentos precisam ser semelhantes ou compatíveis. Consulte a documentação de transmissão para saber o que significa a compatibilidade de formas. O resultado de uma operação tem uma forma que é o resultado da transmissão das duas matrizes de entrada. Nessa variante, as operações entre matrizes de ranks diferentes não são aceitas, a menos que um dos operandos seja um escalar.

Existe uma variante alternativa com suporte de transmissão de dimensões diferentes para o Max:

Max(lhs,rhs, broadcast_dimensions)

Essa variante da operação deve ser usada para operações aritméticas entre matrizes de ranks diferentes (como adicionar uma matriz a um vetor).

O operando broadcast_dimensions adicional é uma fração de números inteiros que especifica as dimensões a serem usadas para transmissão dos operandos. A semântica é descrita em detalhes na página de transmissão.

Para informações sobre o StableHLO, consulte StableHLO: máximo.

Mín.

Consulte também XlaBuilder::Min.

Executa a operação de mínimo com elementos em lhs e rhs.

Min(lhs, rhs)

As formas dos argumentos precisam ser semelhantes ou compatíveis. Consulte a documentação de transmissão para saber o que significa a compatibilidade de formas. O resultado de uma operação tem uma forma que é o resultado da transmissão das duas matrizes de entrada. Nessa variante, as operações entre matrizes de ranks diferentes não são aceitas, a menos que um dos operandos seja um escalar.

Existe uma variante alternativa com suporte de transmissão de dimensões diferentes para Min:

Min(lhs,rhs, broadcast_dimensions)

Essa variante da operação deve ser usada para operações aritméticas entre matrizes de ranks diferentes (como adicionar uma matriz a um vetor).

O operando broadcast_dimensions adicional é uma fração de números inteiros que especifica as dimensões a serem usadas para transmissão dos operandos. A semântica é descrita em detalhes na página de transmissão.

Para informações sobre o StableHLO, consulte StableHLO: mínimo.

Mul

Consulte também XlaBuilder::Mul.

Realiza o produto elemento a elemento de lhs e rhs.

Mul(lhs, rhs)

As formas dos argumentos precisam ser semelhantes ou compatíveis. Consulte a documentação de transmissão para saber o que significa a compatibilidade de formas. O resultado de uma operação tem uma forma que é o resultado da transmissão das duas matrizes de entrada. Nessa variante, as operações entre matrizes de ranks diferentes não são aceitas, a menos que um dos operandos seja um escalar.

Existe uma variante alternativa com suporte de transmissão de dimensões diferentes para Mul:

Mul(lhs,rhs, broadcast_dimensions)

Essa variante da operação deve ser usada para operações aritméticas entre matrizes de ranks diferentes (como adicionar uma matriz a um vetor).

O operando broadcast_dimensions adicional é uma fração de números inteiros que especifica as dimensões a serem usadas para transmissão dos operandos. A semântica é descrita em detalhes na página de transmissão.

Para informações sobre o StableHLO, consulte StableHLO: multiplicação.

Neg

Consulte também XlaBuilder::Neg.

Negação de elemento por elemento x -> -x.

Neg(operand)

Para informações sobre o StableHLO, consulte StableHLO: negação

Não

Consulte também XlaBuilder::Not.

Negação lógica x -> !(x) referente a elemento.

Not(operand)

Para informações sobre o StableHLO, consulte StableHLO - not.

OptimizationBarrier

Consulte também XlaBuilder::OptimizationBarrier.

Impede que qualquer transmissão de otimização mova computações pela barreira.

OptimizationBarrier(operand)

Garante que todas as entradas sejam avaliadas antes de qualquer operador que dependa das saídas da barreira.

Para informações sobre o StableHLO, consulte StableHLO: optimization_barrier.

Ou

Consulte também XlaBuilder::Or.

Executa um OR bit a bit de lhs e rhs .

Or(lhs, rhs)

As formas dos argumentos precisam ser semelhantes ou compatíveis. Consulte a documentação de transmissão para saber o que significa a compatibilidade de formas. O resultado de uma operação tem uma forma que é o resultado da transmissão das duas matrizes de entrada. Nessa variante, as operações entre matrizes de ranks diferentes não são aceitas, a menos que um dos operandos seja um escalar.

Existe uma variante alternativa com suporte de transmissão de dimensões diferentes para "Or":

Or(lhs,rhs, broadcast_dimensions)

Essa variante da operação deve ser usada para operações aritméticas entre matrizes de ranks diferentes (como adicionar uma matriz a um vetor).

O operando broadcast_dimensions adicional é uma fração de números inteiros que especifica as dimensões a serem usadas para transmissão dos operandos. A semântica é descrita em detalhes na página de transmissão.

Para informações sobre o StableHLO, consulte StableHLO - ou.

Outfeed

Consulte também XlaBuilder::Outfeed.

Grava entradas no feed de saída.

Outfeed(operand, shape_with_layout, outfeed_config)

shape_with_layout comunica a forma disposta que queremos extrair.

Para informações sobre o StableHLO, consulte StableHLO: saída.

Pad

Consulte também XlaBuilder::Pad.

Pad(operand, padding_value, padding_config)

Adiciona padding à matriz operand fornecida ao redor dela e entre os elementos com o padding_value especificado. padding_config especifica a quantidade de padding de borda e o padding interno para cada dimensão.

PaddingConfig é um campo repetido de PaddingConfigDimension, que contém três campos para cada dimensão: edge_padding_low, edge_padding_high e interior_padding.

edge_padding_low e edge_padding_high especificam a quantidade de padding adicionada na extremidade inferior (próxima ao índice 0) e na extremidade superior (próxima ao índice mais alto) de cada dimensão, respectivamente. A quantidade de padding de borda pode ser negativa. O valor absoluto do padding negativo indica o número de elementos a serem removidos da dimensão especificada.

interior_padding especifica a quantidade de padding adicionada entre dois elementos em cada dimensão e não pode ser negativa. O padding interno ocorre logicamente antes do padding de borda. Portanto, no caso de padding de borda negativo, os elementos são removidos do operando com padding interno.

Essa operação não faz nada se todos os pares de padding de borda forem (0, 0) e os valores de padding interno forem todos 0. A figura abaixo mostra exemplos de diferentes valores de edge_padding e interior_padding para uma matriz bidimensional.

Para informações sobre o StableHLO, consulte StableHLO: pad.

Parâmetro

Consulte também XlaBuilder::Parameter.

Parameter representa uma entrada de argumento para um cálculo.

PartitionID

Consulte também XlaBuilder::BuildPartitionId.

Produz partition_id do processo atual.

PartitionID(shape)

PartitionID pode aparecer em despejos de HLO, mas não foi criado para ser construído manualmente por usuários finais.

Para informações sobre o StableHLO, consulte StableHLO: partition_id.

PopulationCount

Consulte também XlaBuilder::PopulationCount.

Calcula o número de bits definidos em cada elemento de operand.

PopulationCount(operand)

Para informações sobre o StableHLO, consulte StableHLO: popcnt.

Pow

Consulte também XlaBuilder::Pow.

Realiza a exponenciação elemento a elemento de lhs por rhs.

Pow(lhs, rhs)

As formas dos argumentos precisam ser semelhantes ou compatíveis. Consulte a documentação de transmissão para saber o que significa a compatibilidade de formas. O resultado de uma operação tem uma forma que é o resultado da transmissão das duas matrizes de entrada. Nessa variante, as operações entre matrizes de ranks diferentes não são aceitas, a menos que um dos operandos seja um escalar.

Existe uma variante alternativa com suporte de transmissão de dimensões diferentes para Pow:

Pow(lhs,rhs, broadcast_dimensions)

Essa variante da operação deve ser usada para operações aritméticas entre matrizes de ranks diferentes (como adicionar uma matriz a um vetor).

O operando broadcast_dimensions adicional é uma fração de números inteiros que especifica as dimensões a serem usadas para transmissão dos operandos. A semântica é descrita em detalhes na página de transmissão.

Para informações sobre o StableHLO, consulte StableHLO: energia.

Real

Consulte também XlaBuilder::Real.

Parte real de um elemento de uma forma complexa (ou real). x -> real(x). Se o operando for um tipo de ponto flutuante, Real vai retornar o mesmo valor.

Real(operand)

Para informações sobre o StableHLO, consulte StableHLO: real.

Recv

Consulte também XlaBuilder::Recv.

Recv, RecvWithTokens e RecvToHost são operações que servem como primitivos de comunicação em HLO. Essas operações geralmente aparecem em despejos de HLO como parte de entrada/saída de baixo nível ou transferência entre dispositivos, mas não devem ser construídas manualmente pelos usuários finais.

Recv(shape, handle)

Recebe dados da forma especificada de uma instrução Send em outro cálculo que compartilha o mesmo identificador de canal. Retorna um XlaOp para os dados recebidos.

Para informações sobre o StableHLO, consulte StableHLO: recv.

RecvDone

Consulte também HloInstruction::CreateRecv e HloInstruction::CreateRecvDone.

Assim como Send, a API do cliente da operação Recv representa a comunicação síncrona. No entanto, a instrução é decomposta internamente em duas instruções HLO (Recv e RecvDone) para permitir transferências de dados assíncronas.

Recv(const Shape& shape, int64 channel_id)

Aloca os recursos necessários para receber dados de uma instrução Send com o mesmo channel_id. Retorna um contexto para os recursos alocados, que é usado por uma instrução RecvDone a seguir para aguardar a conclusão da transferência de dados. O contexto é uma tupla de {buffer de recebimento (formato), identificador de solicitação (U32)} e só pode ser usado por uma instrução RecvDone.

Dado um contexto criado por uma instrução Recv, aguarda a conclusão da transferência de dados e retorna os dados recebidos.

Reduzir

Consulte também XlaBuilder::Reduce.

Aplica uma função de redução a uma ou mais matrizes em paralelo.

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

Em que:

• N precisa ser maior ou igual a 1.
• O cálculo precisa ser "aproximadamente" associativo (veja abaixo).
• Todas as matrizes de entrada precisam ter as mesmas dimensões.
• Todos os valores iniciais precisam formar uma identidade em computation.
• Se N = 1, Collate(T) será T.
• Se N > 1, Collate(T_0, ..., T_{N-1}) for uma tupla de elementos N do tipo T.

Essa operação reduz uma ou mais dimensões de cada matriz de entrada em escalares. O número de dimensões de cada matriz retornada é number_of_dimensions(operand) - len(dimensions). A saída da operação é Collate(Q_0, ..., Q_N), em que Q_i é uma matriz do tipo T_i, e as dimensões são descritas abaixo.

Diferentes back-ends podem reassociar o cálculo de redução. Isso pode levar a diferenças numéricas, já que algumas funções de redução, como adição, não são associativas para números de ponto flutuante. No entanto, se o intervalo dos dados for limitado, a adição de ponto flutuante será associativa o suficiente para a maioria dos usos práticos.

Para informações sobre o StableHLO, consulte StableHLO: reduzir.

Exemplos

Ao reduzir uma dimensão em uma única matriz 1D com valores [10, 11, 12, 13] e função de redução f (que é computation), isso pode ser calculado como

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

mas também há muitas outras possibilidades, por exemplo:

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

Confira a seguir um exemplo aproximado de pseudocódigo de como a redução pode ser implementada, usando a soma como o cálculo de redução com um 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]

Confira um exemplo de redução de uma matriz 2D. O formato tem duas dimensões, a dimensão 0 de tamanho 2 e a dimensão 1 de tamanho 3:

Resultados da redução das dimensões 0 ou 1 com uma função "add":

Os dois resultados de redução são matrizes unidimensionais. O diagrama mostra um como coluna e outro como linha apenas para conveniência visual.

Para um exemplo mais complexo, confira uma matriz 3D. O número de dimensões é 3, dimensão 0 de tamanho 4, dimensão 1 de tamanho 2 e dimensão 2 de tamanho 3. Para simplificar, os valores de 1 a 6 são replicados na dimensão 0.

Assim como no exemplo 2D, podemos reduzir apenas uma dimensão. Se reduzirmos a dimensão 0, por exemplo, vamos ter uma matriz bidimensional em que todos os valores da dimensão 0 foram transformados em um escalar:

|  4   8  12 |
| 16  20  24 |

Se reduzirmos a dimensão 2, também vamos receber uma matriz bidimensional em que todos os valores na dimensão 2 foram agrupados em um escalar:

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

A ordem relativa entre as dimensões restantes na entrada é preservada na saída, mas algumas dimensões podem receber novos números, já que o número de dimensões muda.

Também é possível reduzir várias dimensões. A adição de dimensões 0 e 1 reduzidas produz a matriz 1D [20, 28, 36].

Reduzir a matriz 3D em todas as dimensões produz o escalar 84.

Redução variádica

Quando N > 1, a aplicação da função de redução é um pouco mais complexa, já que é aplicada simultaneamente a todas as entradas. Os operandos são fornecidos ao cálculo na seguinte ordem:

• Execução do valor reduzido para o primeiro operando
• …
• Execução de valor reduzido para o Nésimo operando
• Valor de entrada do primeiro operando
• …
• Valor de entrada do N-ésimo operando

Por exemplo, considere a seguinte função de redução, que pode ser usada para calcular o máximo e o argmax de uma matriz unidimensional em paralelo:

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

Para matrizes de entrada 1D V = Float[N], K = Int[N] e valores iniciais I_V = Float, I_K = Int, o resultado f_(N-1) da redução na única dimensão de entrada é equivalente à seguinte aplicação 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))

Aplicar essa redução a uma matriz de valores e uma matriz de índices sequenciais (ou seja, iota) vai iterar em conjunto nas matrizes e retornar uma tupla que contém o valor máximo e o índice correspondente.

ReducePrecision

Consulte também XlaBuilder::ReducePrecision.

Modela o efeito da conversão de valores de ponto flutuante para um formato de menor precisão (como IEEE-FP16) e de volta ao formato original. O número de bits de expoente e mantissa no formato de baixa precisão pode ser especificado arbitrariamente, embora nem todos os tamanhos de bits sejam compatíveis com todas as implementações de hardware.

ReducePrecision(operand, exponent_bits, mantissa_bits)

O resultado é uma matriz do tipo T. Os valores de entrada são arredondados para o valor representável mais próximo com o número de bits de mantissa especificado (usando a semântica "ties to even"), e todos os valores que excedem o intervalo especificado pelo número de bits de expoente são fixados como infinito positivo ou negativo. Os valores NaN são mantidos, mas podem ser convertidos em valores NaN canônicos.

O formato de menor precisão precisa ter pelo menos um bit de expoente (para distinguir um valor zero de um infinito, já que ambos têm uma mantissa zero) e um número não negativo de bits de mantissa. O número de bits de expoente ou mantissa pode exceder o valor correspondente para o tipo T. A parte correspondente da conversão é simplesmente uma operação nula.

Para informações sobre o StableHLO, consulte StableHLO: reduce_precision.

ReduceScatter

Consulte também XlaBuilder::ReduceScatter.

O ReduceScatter é uma operação coletiva que faz um AllReduce e depois dispersa o resultado dividindo-o em shard_count blocos ao longo do scatter_dimension, e a réplica i no grupo de réplicas recebe o fragmento ith.

ReduceScatter(operand, computation, scatter_dimension, shard_count, replica_groups, channel_id, layout, use_global_device_ids)

• Quando operand é uma tupla de matrizes, o reduce-scatter é realizado em cada elemento da tupla.
• replica_groups é uma lista de grupos de réplicas entre os quais a redução é realizada. O ID da réplica atual pode ser recuperado usando ReplicaId. A ordem das réplicas em cada grupo determina a ordem em que o resultado de redução total será disperso. replica_groups precisa estar vazio (nesse caso, todas as réplicas pertencem a um único grupo) ou conter o mesmo número de elementos que o número de réplicas. Quando há mais de um grupo de réplicas, todos precisam ter o mesmo tamanho. Por exemplo, replica_groups = {0, 2}, {1, 3} realiza a redução entre as réplicas 0 e 2, 1 e 3 e, em seguida, dispersa o resultado.
• shard_count é o tamanho de cada grupo de réplicas. Precisamos disso nos casos em que replica_groups estão vazios. Se replica_groups não estiver vazio, shard_count precisará ser igual ao tamanho de cada grupo de réplicas.
• channel_id é usado para comunicação entre módulos: apenas operações reduce-scatter com o mesmo channel_id podem se comunicar entre si.
• layout Consulte xla::shapes para mais informações sobre layouts.
• use_global_device_ids é uma flag especificada pelo usuário. Quando false(padrão), os números em replica_groups são ReplicaId quando true, e replica_groups representa um ID global de (ReplicaID*partition_count + partition_id). Por exemplo:
  • Com duas réplicas e quatro partições,
  • 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)
  • em que cada par é (replica_id, partition_id).

A forma de saída é a forma de entrada com o scatter_dimension shard_count vezes menor. Por exemplo, se houver duas réplicas e o operando tiver o valor [1.0, 2.25] e [3.0, 5.25], respectivamente, nas duas réplicas, o valor de saída dessa operação em que scatter_dim é 0 será [4.0] para a primeira réplica e [7.5] para a segunda.

Para informações sobre o StableHLO, consulte StableHLO: reduce_scatter.

ReduceScatter: exemplo 1 (StableHLO)

No exemplo acima, há duas réplicas participando do ReduceScatter. Em cada réplica, o operando tem o formato f32[2,4]. Uma redução total (soma) é realizada em todas as réplicas, produzindo um valor reduzido de forma f32[2,4] em cada réplica. Esse valor reduzido é dividido em duas partes ao longo da dimensão 1, de modo que cada parte tenha o formato f32[2,2]. Cada réplica no grupo de processos recebe a parte correspondente à sua posição no grupo. Como resultado, a saída em cada réplica tem o formato f32[2,2].

ReduceWindow

Consulte também XlaBuilder::ReduceWindow.

Aplica uma função de redução a todos os elementos em cada janela de uma sequência de matrizes multidimensionais N, produzindo uma única matriz ou uma tupla de N matrizes multidimensionais como saída. Cada matriz de saída tem o mesmo número de elementos que o número de posições válidas da janela. Uma camada de pooling pode ser expressa como um ReduceWindow. Assim como Reduce, o computation aplicado sempre é transmitido pelo init_values no lado esquerdo.

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

Em que:

• N precisa ser maior ou igual a 1.
• Todas as matrizes de entrada precisam ter as mesmas dimensões.
• Se N = 1, Collate(T) será T.
• Se N > 1, Collate(T_0, ..., T_{N-1}) for uma tupla de elementos N do tipo (T0,...T{N-1}).

Para informações sobre o StableHLO, consulte StableHLO: reduce_window.

ReduceWindow - Example 1

A entrada é uma matriz de tamanho [4x6], e window_dimensions e window_stride_dimensions são [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);