Semântica de operação

Mantenha tudo organizado com as coleções Salve e categorize o conteúdo com base nas suas preferências.

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

Observação sobre a nomenclatura: o tipo de dados generalizado com que a XLA lida é uma matriz N-dimensional que contém elementos de algum tipo uniforme, como float de 32 bits. Em toda a documentação, array é usado para denotar uma matriz de dimensão arbitrária. Para sua conveniência, casos especiais têm nomes mais específicos e conhecidos. Por exemplo, um vetor é uma matriz unidimensional, e uma matriz é uma matriz bidimensional.

AfterAll

Consulte também XlaBuilder::AfterAll.

AfterAll recebe um número variável de tokens e produz um único token. Os tokens são tipos primitivos que podem ser encadeados entre operações de efeitos colaterais para forçar a ordenação. O AfterAll pode ser usado como uma mesclagem de tokens para ordenar uma operação após um conjunto de operações.

AfterAll(operands)

AllGather

Consulte também XlaBuilder::AllGather.

Realiza a concatenação em todas as réplicas.

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

• replica_groups é uma lista de grupos de réplicas entre os quais a concatenação é realizada. O ID da réplica para a 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. Isso é necessário nos casos em que replica_groups está vazio.
• channel_id é usado para comunicação entre módulos: apenas operações all-gather com o mesmo channel_id podem se comunicar entre si.

A forma de saída é a forma de entrada com o all_gather_dim multiplicado por shard_count vezes. 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.

AllReduce

Consulte também XlaBuilder::AllReduce.

Realiza uma computação personalizada em todas as réplicas.

AllReduce(operand, computation, replica_group_ids, channel_id)

• Quando operand é uma tupla de matrizes, a redução total é 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 para a réplica atual pode ser recuperado usando ReplicaId. replica_groups precisa estar vazio (caso em que 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.

A forma de saída é igual à forma de 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 da soma será [4.0, 7.75] nas duas réplicas. Se a entrada for uma tupla, a saída também será uma tupla.

O cálculo do resultado de AllReduce requer uma entrada de cada réplica. Portanto, se uma réplica executar um nó AllReduce mais vezes do que outra, a réplica anterior vai esperar para sempre. Como as réplicas estão executando o mesmo programa, não há muitas maneiras de fazer isso, mas é possível quando a condição de um loop while depende de dados da ingestão e os dados que são ingeridos fazem com que o loop while itere mais vezes em uma réplica do que em outra.

AllToAll

Consulte também XlaBuilder::AllToAll.

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

1. A fase de dispersão. Em cada núcleo, o operando é dividido em um número de split_count de blocos ao longo da split_dimensions, e os blocos são espalhados para todos os núcleos. Por exemplo, o bloco ith é enviado para o núcleo ith.
2. Fase de coleta. Cada núcleo concatena os blocos recebidos ao longo do concat_dimension.

As cores participantes podem ser configuradas por:

• replica_groups: cada ReplicaGroup contém uma lista de IDs de réplica que participam da computação. O ID da réplica atual pode ser recuperado usando ReplicaId. AllToAll será aplicado nos subgrupos na ordem especificada. Por exemplo, replica_groups = { {1,2,3}, {4,5,0} } significa que um AllToAll será aplicado nas réplicas {1, 2, 3} e 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.

Pré-requisitos:

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

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

Confira abaixo um exemplo de Alltoall.

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

Neste exemplo, há quatro núcleos participando do Alltoall. Em cada núcleo, o operando é dividido em quatro partes ao longo da dimensão 1. Assim, cada parte tem a forma f32[4,4]. As quatro partes são espalhadas por todas as cores. Em seguida, cada núcleo concatena as partes recebidas na dimensão 0, na ordem dos núcleos 0 a 4. Portanto, a saída em cada núcleo tem a forma f32[16,4].

BatchNormGrad

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

Calcula gradientes de norma de lote.

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

Para cada elemento na dimensão de elementos (feature_index é o índice da dimensão de elementos 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 fórmulas abaixo (considerando uma matriz de quatro dimensões como operand e com índice de dimensão de recurso l, tamanho de 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 mean e variance representam valores de momentos nas dimensões de lote e espacial.

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

BatchNormInference

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

Normaliza uma matriz em dimensões espaciais e de lote.

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

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

BatchNormInference é equivalente a chamar BatchNormTraining sem calcular mean e variance para cada lote. Ele usa a entrada 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 a mesma forma da entrada operand.

BatchNormTraining

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

Normaliza uma matriz em dimensões espaciais e de lote.

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

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

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

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

• Calcula a variação de 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 desloca: \(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:

O batch_mean e o batch_var são momentos calculados nas dimensões de lote e espacial usando as fórmulas acima.

BitcastConvertType

Consulte também XlaBuilder::BitcastConvertType.

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

BitcastConvertType(operand, new_element_type)

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

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

Conversão de bitcast para tipo primitivo de largura diferente

A instrução HLO BitcastConvert oferece suporte ao 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', a forma de saída recebe uma nova dimensão 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 bit não são por 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. Ou seja, 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}.

As novas dimensões são indexadas em cópias do operando, ou seja,

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

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

BroadcastInDim

Consulte também XlaBuilder::BroadcastInDim.

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

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 o 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 do formato de destino, ou seja, a i-ésima dimensão do operando é mapeada para a broadcast_dimension[i] do formato de saída. As dimensões de operand precisam ter o tamanho 1 ou o mesmo tamanho da dimensão na forma de saída para a qual elas são mapeadas. As dimensões restantes são preenchidas com dimensões de tamanho 1. A transmissão de dimensão degenerada é transmitida ao longo dessas dimensões degeneradas 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 uma computação com os argumentos fornecidos.

Call(computation, args...)

A aritmética e os tipos de args precisam corresponder aos parâmetros de computation. É permitido não ter args.

CompositeCall

Consulte também XlaBuilder::CompositeCall.

Encapsula uma operação composta por outras operações StableHLO, usando 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 alterar a semântica do programa. Nos casos em que a inserçã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 com o prefixo 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"
}

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

Cholesky

Consulte também XlaBuilder::Cholesky.

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

Cholesky(a, lower)

Se lower for true, calcule matrizes triangulares inferiores l de modo que $a = l . l^T$. Se lower for false, calcule as matrizes superiores triangulares u de modo 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, a será tratado como um lote de matrizes, em que todas, exceto as duas dimensões menores, são dimensões de lote.

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

Limitar

Consulte também XlaBuilder::Clamp.

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

As três matrizes precisam ter a mesma forma. 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};

Recolher

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

Reúne as dimensões de uma matriz em uma.

Collapse(operand, dimensions)

A opção "Recolher" 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 estar em ordem (números de dimensão de baixo para alto), subconjunto 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 que elas substituem, com o novo tamanho da dimensão igual ao produto dos tamanhos das dimensões originais. O número de dimensão mais baixo em dimensions é a dimensão variável mais lenta (mais importante) no conjunto de loops que une essas dimensões, e o número de dimensão mais alto é a variação mais rápida (menos importante). Consulte o operador tf.reshape se for necessário ordenar mais geralmente a ordenação de fechamento.

Por exemplo, 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} };

CollectivePermute

Consulte também XlaBuilder::CollectivePermute.

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

CollectivePermute(operand, source_target_pairs)

Observe que há as seguintes restrições na source_target_pair:

• Dois pares não podem ter o mesmo ID de réplica de destino nem o mesmo ID de réplica 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 por 0s com a mesma forma da entrada.

Concatenate

Consulte também XlaBuilder::ConcatInDim.

A concatenação compõe uma matriz de vários operandos de matriz. A matriz tem o mesmo número de dimensões que cada um dos operandos da matriz de entrada (que precisam ter o mesmo número de dimensões) 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 ocorre porque a XLA não oferece suporte a matrizes "ragged". Além disso, valores de dimensão 0 não podem ser concatenados, porque é impossível nomear a dimensão em que a concatenação ocorre.

Exemplo unidimensional:

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

Exemplo bidimensional:

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

Diagrama:

Condicional

Consulte também XlaBuilder::Conditional.

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

Executa true_computation se pred for true, false_computation se pred 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 receber 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 pred.

Conditional(branch_index, branch_computations, branch_operands)

Executa branch_computations[branch_index] e retorna o resultado. Se branch_index for um S32 que é < 0 ou >= N, então 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.

Conv (convolução)

Consulte também XlaBuilder::Conv.

Como ConvWithGeneralPadding, mas o padding é especificado de forma abreviada como SAME ou VALID. O preenchimento idêntico preenche a entrada (lhs) com zeros para que a saída tenha a mesma forma da entrada quando não considera o passo. VALID padding significa simplesmente que não há padding.

ConvWithGeneralPadding (convolução)

Consulte também XlaBuilder::ConvWithGeneralPadding.

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 uma computação é realizada para cada posição possível da janela.

Vamos supor que n seja o número de dimensões espaciais. O argumento lhs é uma matriz dimensional (n+2) que descreve a área da base. Isso é chamado de entrada, mas o lado direito também é 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 em que 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 de (n+2) dimensões 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 o passo da janela de convolução nas dimensões espaciais. Por exemplo, se o passo na primeira dimensão espacial for 3, a janela só poderá ser colocada em coordenadas em que o primeiro índice espacial for divisível por 3.

O argumento padding especifica a quantidade de padding 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 elemento. O padding baixo é aplicado na direção de índices mais baixos, enquanto o padding alto é aplicado na direção de índices mais altos. Por exemplo, se padding[1] for (2,3), haverá um padding de dois zeros à esquerda e de três zeros à direita na segunda dimensão espacial. Usar o 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, então d-1 furos serão colocados implicitamente entre cada uma das entradas nessa dimensão, aumentando o tamanho da matriz. Os buracos são preenchidos com um valor sem operação, que para a convolução significa zeros.

A dilatação do lado direito também é chamada de convolução de Atrous. Para mais detalhes, consulte tf.nn.atrous_conv2d. A dilatação da LH 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 entrada e da dimensão de recurso de saída. Se feature_group_count for maior que 1, isso significa que conceitualmente a dimensão de recursos de entrada e saída e a dimensão de recursos de saída rhs são divididas igualmente em muitos grupos feature_group_count, cada um deles com uma subseqência consecutiva de recursos. A dimensão do recurso de entrada de rhs precisa ser igual à dimensão do recurso de entrada lhs dividida por feature_group_count. Assim, ela já tem o tamanho de um grupo de recursos de entrada. Os grupos i 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 do elemento de saída.

Para a convolução de profundidade, o argumento feature_group_count seria definido como a dimensão do elemento de entrada, e o filtro seria remodelado 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 retropropagação. batch_group_count precisa ser um divisor do tamanho da dimensão de lote lhs (entrada). Se batch_group_count for maior que 1, isso significa que a dimensão do lote de saída precisa ter o tamanho input batch / batch_group_count. O batch_group_count precisa ser um divisor do tamanho do elemento de saída.

A forma de 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 no lado esquerdo.
• z: tem o mesmo tamanho que output-z no kernel (rhs).
• spatial_dims: um valor para cada posicionamento válido da janela convolucional.

A figura acima mostra como o campo batch_group_count funciona. Na prática, cortamos cada lote de lhs em grupos batch_group_count e fazemos o mesmo para os recursos de saída. Em seguida, para cada um desses grupos, fazemos convoluções em pares e concatenamos a saída ao longo da dimensão do elemento 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 de convolução são determinadas pelos passos e pelo tamanho da área de base após o preenchimento.

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, tirada da área de base, em que cada ponto 2D é associado a um vetor 1D, de modo que temos uma caixa 3D. Do kernel de convolução, como corrigimos a coordenada de saída z, também temos uma caixa 3D. As duas caixas têm as mesmas dimensões, então podemos usar a soma dos produtos de elementos entre as duas caixas (semelhante a um produto de pontos). Esse é o valor de saída.

Se output-z for, por exemplo, 5, então cada posição da janela produz 5 valores na saída na dimensão z da saída. Esses valores diferem na parte do kernel convolucional usada. Há uma caixa 3D separada de valores usados para cada coordenada output-z. 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 preenchimento e passo:

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

ConvertElementType

Consulte também XlaBuilder::ConvertElementType.

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

ConvertElementType(operand, new_element_type)

As dimensões do operando e da forma de destino precisam ser correspondentes. 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 normalização de int para float, como arredondar 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}

CrossReplicaSum

Realiza AllReduce com um cálculo de soma.

CustomCall

Consulte também XlaBuilder::CustomCall.

Chamar uma função fornecida pelo usuário em uma computação.

CustomCall(target_name, args..., shape)

A assinatura da função é a mesma, independentemente da aritmética ou do tipo de argumentos:

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

Por exemplo, se a CustomCall for usada da seguinte maneira:

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

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

Confira um exemplo de implementação de myfunc:

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

A função fornecida pelo usuário não pode ter efeitos colaterais, e a execução dela precisa ser idempotente.

Dot

Consulte também XlaBuilder::Dot.

Dot(lhs, rhs)

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 apenas uma dimensão) e na primeira dimensão de rhs. Essas são as dimensões "contratuais". As dimensões contratadas 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.

DotGeneral

Consulte também XlaBuilder::DotGeneral.

DotGeneral(lhs, rhs, dimension_numbers)

Semelhante a "Ponto", mas permite que os números de dimensão de contrato e lote sejam especificados para lhs e rhs.

O DotGeneral realiza a soma dos produtos em dimensões de contrataçã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 contrato:

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

O número da dimensão resultante começa com a dimensão de lote, depois a dimensão lhs não de lote/contratação e, por fim, a dimensão rhs não de lote/contratação.

DynamicSlice

Consulte também XlaBuilder::DynamicSlice.

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

DynamicSlice(operand, start_indices, size_indices)

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

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

Isso garante que a fatia extraída esteja sempre dentro dos limites em relação à matriz de operando. Se o slice estiver dentro dos limites antes da transformação ser aplicada, 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}) produces:
{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}) produces:
{ { 7.0,  8.0},
{10.0, 11.0} }

DynamicUpdateSlice

Consulte também XlaBuilder::DynamicUpdateSlice.

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

DynamicUpdateSlice(operand, update, start_indices)

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

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 o slice estiver dentro dos limites antes da transformação ser aplicada, 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) produces:
{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) produces:
{ {0.0,  1.0,  2.0},
{3.0, 12.0, 13.0},
{6.0, 14.0, 15.0},
{9.0, 16.0, 17.0} }

Operações aritméticas binárias elementar

Consulte também XlaBuilder::Add.

Um conjunto de operações aritméticas binárias por elemento é aceito.

Op(lhs, rhs)

Onde Op é uma das seguintes operações: Add (adição), Sub(subtração), Mul (multiplicação), Div (divisão), Pow (potência), Rem (resto), Max (máximo), Min (mínimo), And (E lógico), Or (OU lógico), Xor (XOU lógico), ShiftLeft (deslocamento para a esquerda), ShiftRightArithmetic (deslocamento para a direita aritmética), ShiftRightLogical (deslocamento para a direita lógico), Atan2 (arcotangente de dois argumentos) ou Complex (combina partes reais e imaginárias em um número complexo).

As formas dos argumentos precisam ser semelhantes ou compatíveis. Consulte a documentação de transmissão sobre o que significa que as formas sejam compatíveis. 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 diferentes níveis não são aceitas, a menos que um dos operandos seja um escalar.

Quando Op é Rem, o sinal do resultado é retirado do dividendo, e o valor absoluto do resultado é sempre menor que o valor absoluto do divisor.

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

Uma variante alternativa com suporte a transmissão de diferentes dimensões existe para essas operações:

Op(lhs, rhs, broadcast_dimensions)

Em que Op é o mesmo que acima. Essa variante da operação precisa ser usada para operações aritméticas entre matrizes de diferentes níveis (como adicionar uma matriz a um vetor).

O operando broadcast_dimensions adicional é uma fatia de números inteiros usada para aumentar o número de dimensões do operando de menor dimensão até o número de dimensões do operando de maior dimensão. broadcast_dimensions mapeia as dimensões da forma de menor dimensão para as dimensões da forma de maior dimensão. As dimensões não mapeadas da forma expandida são preenchidas com dimensões de tamanho um. A transmissão de dimensão degenerada transmite as formas ao longo dessas dimensões degeneradas para igualar as formas dos dois operandos. A semântica é descrita em detalhes na página de transmissão.

Operações de comparação elementar

Consulte também XlaBuilder::Eq.

Um conjunto de operações de comparação binária elementar padrão é aceito. Observe que a semântica de comparação de ponto flutuante IEEE 754 padrão é aplicada ao comparar tipos de ponto flutuante.

Op(lhs, rhs)

Em que Op é Eq (igual a), Ne (diferente de), Ge (maior ou igual a), Gt (maior que), Le (menor ou igual a) ou Lt (menor que). Outro conjunto de operadores, EqTotalOrder, NeTotalOrder, GeTotalOrder, GtTotalOrder, LeTotalOrder e LtTotalOrder, oferece as mesmas funcionalidades, exceto pelo fato de que eles também oferecem suporte a uma ordem total sobre os números de ponto flutuante, forçando -NaN < -Inf < -Finite < -0 < +0 < +Finite < +Inf < +NaN.

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

Uma variante alternativa com suporte a transmissão de diferentes dimensões existe para essas operações:

Op(lhs, rhs, broadcast_dimensions)

Em que Op é o mesmo que acima. Essa variante da operação precisa ser usada para operações de comparação entre matrizes de diferentes níveis (como adicionar uma matriz a um vetor).

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

Funções unárias elementar

O XlaBuilder oferece suporte a estas funções unárias por elemento:

Abs(operand) Elemento x -> |x| absoluto.

Cbrt(operand) Operação de raiz cúbica elementar x -> cbrt(x).

Ceil(operand) Elemento de x -> ⌈x⌉.

Clz(operand) Conta os zeros à esquerda de cada elemento.

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

Erf(operand) Função de erro elementar x -> erf(x) em que

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

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

Expm1(operand) Exponencial natural elementar menos um x -> e^x - 1.

Floor(operand) Elemento de x -> ⌊x⌋.

Imag(operand) Parte imaginária elementar de uma forma complexa (ou real). x -> imag(x). Se o operando for do tipo ponto flutuante, retornará 0.

IsFinite(operand) 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 a mesma forma da entrada, em que cada elemento é true se e somente se o elemento de entrada correspondente for finito.

Log(operand) Logaritmo natural elementar x -> ln(x).

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

Logistic(operand) Cálculo da função logística elementar x -> logistic(x).

Neg(operand) Negação elementar x -> -x.

Not(operand) Negação lógica elementar x -> !(x).

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

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

Round(operand) Arredondamento por elemento, empates longe de zero.

RoundNearestEven(operand) Arredondamento por elemento, vinculação ao par mais próximo.

Rsqrt(operand) Reciprocidade elementar da operação de raiz quadrada x -> 1.0 / sqrt(x).

Sign(operand) Operação de sinal elementar x -> sgn(x) em que

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

usando o operador de comparação do tipo de elemento operand.

Sin(operand) Função seno por elemento x -> sin(x).

Sqrt(operand) Operação de raiz quadrada elementar x -> sqrt(x).

Tan(operand) Tangente elementar x -> tan(x).

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

A função é aplicada a cada elemento na matriz operand, resultando em uma matriz com a mesma forma. É permitido que operand seja um escalar (0 dimensões).

Ft

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

Consulte também XlaBuilder::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->complex e complex->real, a transformação do eixo mais interno é realmente realizada primeiro (RFFT; o último para IRFFT), e é por isso que o eixo mais interno é o que muda de tamanho. Outras transformações de eixo serão complexas->complexas.

Detalhes da implementação

O FFT da CPU é suportado pelo TensorFFT do Eigen. O FFT da GPU usa cuFFT.

Gather

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

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, offset_dims, collapsed_slice_dims, slice_sizes, start_index_map)

Por conveniência, marcamos 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 dimensões batch_dims.size + offset_dims.size.

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, consideraremos implicitamente que start_indices tem uma dimensão 1 final (ou seja, se start_indices tiver a forma [6,7] e index_vector_dim for 2, consideraremos implicitamente que a forma de start_indices é [6,7,1]).

Os limites da matriz de saída na dimensão i são calculados da seguinte maneira:

1. Se i estiver presente em batch_dims (ou seja, for igual a batch_dims[k] para alguns k), escolheremos os limites de dimensão correspondentes de start_indices.shape, pulando index_vector_dim (ou seja, escolha 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 algumas k), vamos escolher o limite correspondente de slice_sizes após 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 de operando In correspondente a um determinado índice de saída Out é calculado da seguinte maneira:

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

2. Crie um índice inicial, Sin, em operand usando S espalhando 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, espalhando 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 elementar.

remapped_offset_dims é uma função monótona com domínio [0, offset_dims.size] e intervalo [0, operand.rank] \ collapsed_slice_dims. 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, a XLA poderá presumir que start_indices forem classificados (em ordem crescente, depois de espalhar os valores de acordo com start_index_map) pelo usuário. Caso contrário, a semântica é definida pela implementação.

Descrição informal e exemplos

Informalmente, cada índice Out na matriz de saída corresponde a um elemento E na matriz de operando, calculado da seguinte maneira:

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

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

• Cortamos dinamicamente uma fatia com o tamanho slice_sizes usando o índice de início completo.

• Reformulamos o recorte encerrando as dimensões collapsed_slice_dims. Como todas as dimensões de fatia recolhidas precisam ter um limite de 1, essa reformulação é 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 de forma fundamental, mas tornam a representação visual mais complicada.

Para entender como tudo isso se encaixa, vamos analisar um exemplo que reúne cinco fatias da 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 da 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 representado como uma transformação de índice que recebe [G,O0,O1], um índice na forma 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₁, que por sua vez decide os limites do corte.

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

Os índices de coleta podem ser multidimensionais. Por exemplo, uma versão mais geral do exemplo acima usando uma matriz "gather indices" do formato [4,5,2] transformaria os índices desta forma:

Novamente, isso funciona como uma fatia dinâmica de lote G0 e G1 como as dimensões do lote. O tamanho da fatia ainda é [8,6].

A operação de coleta no XLA generaliza a semântica informal descrita acima da seguinte maneira:

1. É possível configurar quais dimensões na forma de saída são as dimensões de deslocamento (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 dimensões de deslocamento.

2. O número de dimensões de deslocamento de saída presentes explicitamente 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 fatia de 1. Como eles têm um tamanho de fatia de 1, o único índice válido para eles é 0, e a eliminação deles não introduz 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 precisa 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 deslocamento de saída com o valor O0. No entanto, antes de serem usados como índices no array de entrada, eles são expandidos de acordo com o "mapeamento de índice de coleta" (start_index_map na descrição formal) e o "mapeamento de deslocamento" (remapped_offset_dims na descrição formal) em [X,0] e [0,O0], respectivamente, somando-se a [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 de tf.gather_nd.

O slice_sizes para este caso é [1,11]. Intuitivamente, 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 formato de matriz.

GetDimensionSize(operand, dimension)

SetDimensionSize

Consulte também XlaBuilder::SetDimensionSize.

Define o tamanho dinâmico da dimensão especificada da XlaOp. O operando precisa ter formato de matriz.

SetDimensionSize(operand, size, dimension)

Transmite o operando como resultado, com a dimensão dinâmica rastreada pelo compilador.

Os valores adicionados serão ignorados pelas operações de redução downstream.

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

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

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

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

GetTupleElement

Consulte também XlaBuilder::GetTupleElement.

Indexa em uma tupla com um valor constante no 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.

Infeed

Consulte também XlaBuilder::Infeed.

Infeed(shape)

Lê um único item de dados da interface de streaming de Infeed implícita do dispositivo, interpretando os dados como a forma e o layout fornecidos e retornando um XlaOp dos dados. Várias operações de Infeed são permitidas em uma computação, mas é necessário que haja uma ordem total entre as operações de Infeed. Por exemplo, dois Infeeds no código abaixo têm uma ordem total, já que 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 formas de tupla aninhadas. Para uma forma de tupla vazia, a operação de Infeed é efetivamente uma operação sem efeito e prossegue sem ler nenhum dado da Infeed do dispositivo.

Iota

Consulte também XlaBuilder::Iota.

Iota(shape, iota_dimension)

Cria um literal constante no dispositivo em vez de uma transferência de host potencialmente grande. Cria uma matriz com a forma especificada e armazena valores a partir de zero e incrementa 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 ]]

Mapa

Consulte também XlaBuilder::Map.

Map(operands..., computation)

Aplica uma função escalar às matrizes operands especificadas, produzindo uma matriz das 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 é uma computação arbitrária com a restrição de que ela tem N entradas de tipo escalar T e uma única saída com o tipo S. A saída tem as mesmas dimensões dos 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.

OptimizationBarrier

Bloqueia qualquer passagem de otimização para mover cálculos pela barreira.

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

Almofada

Consulte também XlaBuilder::Pad.

Pad(operand, padding_value, padding_config)

Expande a matriz operand especificada adicionando padding ao redor da matriz e entre os elementos da matriz 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 preenchimento adicionado no extremo inferior (ao lado do índice 0) e no extremo superior (ao lado do í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. Ela 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 será realizada se os pares de preenchimento de borda forem todos (0, 0) e os valores de preenchimento interno forem todos 0. A figura abaixo mostra exemplos de diferentes valores de edge_padding e interior_padding para uma matriz bidimensional.

Recv

Consulte também XlaBuilder::Recv.

Recv(shape, channel_handle)

Recebe dados da forma especificada de uma instrução Send em outra computação que compartilha o mesmo identificador de canal. Retorna uma XlaOp para os dados recebidos.

A API 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. Consulte também HloInstruction::CreateRecv e HloInstruction::CreateRecvDone.

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 seguinte para aguardar a conclusão da transferência de dados. O contexto é uma tupla de {receive buffer (shape), request identifier (U32)} e só pode ser usado por uma instrução RecvDone.

RecvDone(HloInstruction context)

Dado um contexto criado por uma instrução Recv, aguarda a transferência de dados ser concluída 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)

Em que:

• N precisa ser maior ou igual a 1.
• A computação precisa ser "aproximadamente" associativa (consulte 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) é T.
• Se N > 1, Collate(T_0, ..., T_{N-1}) é 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, cujas dimensões são descritas abaixo.

É possível associar novamente o cálculo de redução a back-ends diferentes. Isso pode levar a diferenças numéricas, já que algumas funções de redução, como a adição, não são associativas para números flutuantes. No entanto, se o intervalo de dados for limitado, a adição de ponto flutuante será suficientemente próxima para ser associativa na maioria dos usos práticos.

Exemplos

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

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

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

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

O exemplo a seguir é um pseudocódigo aproximado de como a redução pode ser implementada, usando a soma como a computação 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 (matriz) bidimensional. A forma 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 "adicionar":

Os dois resultados de redução são matrizes 1D. O diagrama mostra uma como coluna e outra como linha apenas para conveniência visual.

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

Da mesma forma que 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 dobrados 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 da dimensão 2 foram dobrados 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. Adicionar as dimensões de redução 0 e 1 produz a matriz 1D [20, 28, 36].

A redução da matriz 3D em todas as dimensões produz o escalar 84.

Variadic reduce

Quando N > 1, a aplicação da função de redução é um pouco mais complexa, porque é aplicada simultaneamente a todas as entradas. Os operandos são fornecidos à computação na seguinte ordem:

• Executando o valor reduzido para o primeiro operando
• …
• Executando o valor reduzido para o N-ésimo operando
• Valor de entrada do primeiro operando
• …
• Valor de entrada para o 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 de 1D 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 unidimensionais V = Float[N], K = Int[N] e valores de inicialização I_V = Float, I_K = Int, o resultado f_(N-1) da redução em toda a dimensão de entrada é equivalente ao seguinte aplicativo recursivo:

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 seqüenciais (por exemplo, iota) vai iterar em conjunto as 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 precisão menor (como IEEE-FP16) e de volta para o formato original. O número de bits de expoente e mantissa no formato de precisão inferior pode ser especificado arbitrariamente, embora nem todos os tamanhos de bit tenham suporte em todas as implementações de hardware.

ReducePrecision(operand, mantissa_bits, exponent_bits)

O resultado é uma matriz do tipo T. Os valores de entrada são arredondados para o valor mais próximo que pode ser representado com o número de bits de mantissa (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 limitados a infinito positivo ou negativo. Os valores de NaN são retidos, embora possam ser convertidos em valores canônicos de NaN.

O formato de menor precisão precisa ter pelo menos um bit exponencial para distinguir um valor zero de um infinito, já que ambos têm uma mantissa zero, e precisa ter 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.

ReduceScatter

Consulte também XlaBuilder::ReduceScatter.

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

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

• Quando operand é uma tupla de matrizes, o espalhamento de redução é 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 para a 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á espalhado. 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} executa a redução entre as réplicas 0 e 2 e 1 e 3 e depois dispersa o resultado.
• shard_count é o tamanho de cada grupo de réplicas. Isso é necessário nos casos em que replica_groups está vazio. Se replica_groups não estiver vazio, shard_count precisa ser igual ao tamanho de cada grupo de réplica.
• channel_id é usado para comunicação entre módulos: apenas operações reduce-scatter com o mesmo channel_id podem se comunicar entre si.

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.

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 N matrizes multidimensionais, produzindo uma única tupla ou 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 agregação pode ser expressa como uma ReduceWindow. Semelhante a Reduce, o computation aplicado é sempre transmitido ao 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) é T.
• Se N > 1, Collate(T_0, ..., T_{N-1}) é uma tupla de elementos N do tipo (T0,...T{N-1}).

O código e a figura abaixo mostram um exemplo de uso de ReduceWindow. 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);

O passo de 1 em uma dimensão especifica que a posição de uma janela na dimensão está a 1 elemento de distância da janela adjacente. Para especificar que nenhuma janela se sobrepõe, window_stride_dimensions precisa ser igual a window_dimensions. A figura abaixo ilustra o uso de dois valores de passo diferentes. O padding é aplicado a cada dimensão da entrada, e os cálculos são os mesmos como se a entrada tivesse vindo com as dimensões que ela tem após o padding.

Para um exemplo de preenchimento não trivial, considere calcular o mínimo da janela de redução (o valor inicial é MAX_FLOAT) com a dimensão 3 e o passo 2 na matriz de entrada [10000, 1000, 100, 10, 1]. O preenchimento kValid calcula mínimos em duas janelas válidas: [10000, 1000, 100] e [100, 10, 1], resultando na saída [100, 1]. O preenchimento kSame preenche primeiro a matriz para que a forma após a janela de redução seja a mesma que a entrada para o passo um, adicionando elementos iniciais em ambos os lados, recebendo [MAX_VALUE, 10000, 1000, 100, 10, 1, MAX_VALUE]. A execução de reduce-window sobre a matriz preenchida opera em três janelas [MAX_VALUE, 10000, 1000], [1000, 100, 10], [10, 1, MAX_VALUE] e produz [1000, 10, 1].

A ordem de avaliação da função de redução é arbitrária e pode ser não determinística. Portanto, a função de redução não deve ser muito sensível à rea associação. Consulte a discussão sobre associatividade no contexto de Reduce para mais detalhes.

ReplicaId

Consulte também XlaBuilder::ReplicaId.

Retorna o ID exclusivo (escalar U32) da réplica.

ReplicaId()

O ID exclusivo de cada réplica é um número inteiro não assinado no intervalo [0, N), em que N é o número de réplicas. Como todas as réplicas estão executando o mesmo programa, uma chamada ReplicaId() no programa vai retornar um valor diferente em cada réplica.

Remodelação

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

Remodela as dimensões de uma matriz em uma nova configuração.

Reshape(operand, dimensions)

Conceitualmente, a reformulação primeiro aplana uma matriz em um vetor unidimensional de valores de dados e, em seguida, refina esse vetor em uma nova forma. Os argumentos de entrada são uma matriz arbitrária do tipo T, um vetor constante de tempo de compilação de índices de dimensão e um vetor constante de tempo de compilação de tamanhos de dimensão para o resultado. O vetor dimensions determina o tamanho da matriz de saída. O valor no índice 0 em dimensions é o tamanho da dimensão 0, o valor no índice 1 é o tamanho da dimensão 1 e assim por diante. O produto das dimensões dimensions precisa ser igual ao produto dos tamanhos de dimensão do operando. Ao refinar a matriz reduzida para a matriz multidimensional definida por dimensions, as dimensões em dimensions são ordenadas da variação mais lenta (mais importante) para a mais rápida (mais secundária).

Por exemplo, 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} } };

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

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

Como um caso especial, a transformação pode transformar uma matriz de um único elemento em um escalar e vice-versa. Por exemplo,

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

Rev (reversa)

Consulte também XlaBuilder::Rev.

Rev(operand, dimensions)

Inverte a ordem dos elementos na matriz operand ao longo do dimensions especificado, gerando uma matriz de saída com a mesma forma. Cada elemento da matriz de operando em um índice multidimensional é armazenado na matriz de saída em um índice transformado. O índice multidimensional é transformado revertendo o índice em cada dimensão a ser revertida. Ou seja, se uma dimensão de tamanho N for uma das dimensões de reversão, o índice i dela será transformado em N - 1 - i.

Um uso da operação Rev é reverter a matriz de peso de convolução ao longo das duas dimensões da janela durante a computação de gradiente em redes neurais.

RngNormal

Consulte também XlaBuilder::RngNormal.

Constrói uma saída de uma forma específica com números aleatórios gerados seguindo a \(N(\mu, \sigma)\) distribuição normal. Os parâmetros \(\mu\) e \(\sigma\)e o formato de saída precisam ter um tipo elementar de ponto flutuante. Além disso, os parâmetros precisam ter valor escalar.

RngNormal(mu, sigma, shape)

RngUniform

Consulte também XlaBuilder::RngUniform.

Constrói uma saída de uma determinada forma com números aleatórios gerados seguindo a distribuição uniforme no intervalo \([a,b)\). Os parâmetros e o tipo de elemento de saída precisam ser um tipo booleano, um tipo integral ou um tipo de ponto flutuante, e os tipos precisam ser consistentes. No momento, os back-ends de CPU e GPU oferecem suporte apenas a F64, F32, F16, BF16, S64, U64, S32 e U32. Além disso, os parâmetros precisam ter valor escalar. Se \(b <= a\) , o resultado é definido pela implementação.

RngUniform(a, b, shape)

RngBitGenerator

Gera uma saída com uma determinada forma preenchida com bits aleatórios uniformes usando o algoritmo especificado (ou padrão do back-end) e retorna um estado atualizado (com a mesma forma que o estado inicial) e os dados aleatórios gerados.

O estado inicial é o estado inicial da geração de números aleatórios atual. Ele e a forma e os valores válidos necessários dependem do algoritmo usado.

A saída é garantida como uma função determinística do estado inicial, mas não é garantido que seja determinística entre back-ends e diferentes versões do compilador.

RngBitGenerator(algorithm, key, shape)

Valores disponíveis para algorithm:

• rng_default: algoritmo específico do back-end com requisitos de forma específicos do back-end.

• rng_three_fry: algoritmo PRNG baseado em contador ThreeFry. A forma initial_state é u64[2] com valores arbitrários. Salmon et al. SC 2011. Números aleatórios paralelos: é só seguir as instruções.

• rng_philox: algoritmo Philox para gerar números aleatórios em paralelo. A forma initial_state é u64[3] com valores arbitrários. Salmon et al. SC 2011. Números aleatórios paralelos: é só seguir as instruções.

Dispersão

A operação de dispersão XLA gera uma sequência de resultados, que são os valores da matriz de entrada operands, com várias fatias (em índices especificados por scatter_indices) atualizadas com a sequência de valores em updates usando update_computation.

Consulte também XlaBuilder::Scatter.

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

Em que:

• N precisa ser maior ou igual a 1.
• operands[0], ..., operands[N-1] precisam ter as mesmas dimensões.
• updates[0], ..., updates[N-1] precisam ter as mesmas dimensões.
• Se N = 1, Collate(T) é T.
• Se N > 1, Collate(T_0, ..., T_N) é uma tupla de elementos N do tipo T.

Se index_vector_dim for igual a scatter_indices.rank, consideraremos implicitamente que scatter_indices tem uma dimensão 1 final.

Definimos update_scatter_dims do tipo ArraySlice<int64> como o conjunto de dimensões no formato updates que não estão em update_window_dims, em ordem ascendente.

Os argumentos de dispersão precisam seguir estas restrições:

• Cada matriz updates precisa ter dimensões update_window_dims.size + scatter_indices.rank - 1.

• Os limites da dimensão i em cada matriz updates precisam estar de acordo com o seguinte:

  • Se i estiver presente em update_window_dims (ou seja, igual a update_window_dims[k] para alguns k), o limite da dimensão i em updates não poderá exceder o limite correspondente de operand após considerar o inserted_window_dims (ou seja, adjusted_window_bounds[k], em que adjusted_window_bounds contém os limites de operand com os limites nos índices inserted_window_dims removidos).
  • Se i estiver presente em update_scatter_dims (ou seja, igual a update_scatter_dims[k] para alguns k), o limite da dimensão i em updates precisará ser igual ao limite correspondente de scatter_indices, ignorando index_vector_dim (ou seja, scatter_indices.shape.dims[k], se k < index_vector_dim e scatter_indices.shape.dims[k+1] caso contrário).

• update_window_dims precisa estar em ordem crescente, não pode ter números de dimensão repetidos e precisa estar no intervalo [0, updates.rank).

• inserted_window_dims precisa estar em ordem crescente, não pode ter números de dimensão repetidos e precisa estar no intervalo [0, operand.rank).

• operand.rank precisa ser igual à soma de update_window_dims.size e inserted_window_dims.size.

• scatter_dims_to_operand_dims.size precisa ser igual a scatter_indices.shape.dims[index_vector_dim], e os valores precisam estar no intervalo [0, operand.rank).

Para um determinado índice U em cada matriz updates, o índice I correspondente na matriz operands correspondente em que essa atualização precisa ser aplicada é calculado da seguinte maneira:

1. G = { U[k] para k em update_scatter_dims }. Use G para procurar um vetor de índice S na matriz scatter_indices, de modo que S[i] = scatter_indices[Combine(G, i)] em que Combine(A, b) insere b nas posições index_vector_dim em A.
2. Crie um índice Sin em operand usando S espalhando S usando o mapa scatter_dims_to_operand_dims. Mais formalmente:
   1. Sin[scatter_dims_to_operand_dims[k]] = S[k] se k < scatter_dims_to_operand_dims.size.
   2. Sin[_] = 0 caso contrário.
3. Crie um índice Win em cada matriz operands espalhando os índices em update_window_dims em U de acordo com inserted_window_dims. Mais formalmente:
   1. Win[window_dims_to_operand_dims(k)] = U[k] se k estiver em update_window_dims, em que window_dims_to_operand_dims é a função monotonicamente crescente com domínio [0, update_window_dims.size] e intervalo [0, operand.rank] \ inserted_window_dims. Por exemplo, se update_window_dims.size for 4, operand.rank for 6 e inserted_window_dims for {0, 2}, window_dims_to_operand_dims será {0→1, 1→3, 2→4, 3→5}.
   2. Win[_] = 0 caso contrário.
4. I é Win + Sin, em que + é a adição elementar.

Em resumo, a operação de dispersão pode ser definida da seguinte maneira.

• Inicializar output com operands, ou seja, para todos os índices J, para todos os índices O na matriz operands[J]:
  output[J][O] = operands[J][O]
• Para cada índice U na matriz updates[J] e o índice correspondente O na matriz operand[J], se O for um índice válido para output:
  (output[0][O], ..., output[N-1][O]) =update_computation(output[0][O], ..., ,output[N-1][O],updates[0][U], ...,updates[N-1][U])

A ordem em que as atualizações são aplicadas não é determinística. Portanto, quando vários índices em updates se referem ao mesmo índice em operands, o valor correspondente em output não é determinístico.

O primeiro parâmetro transmitido para update_computation será sempre o valor atual da matriz output, e o segundo parâmetro será sempre o valor da matriz updates. Isso é importante especificamente para casos em que o update_computation não é comutativo.

Se indices_are_sorted for definido como verdadeiro, a XLA poderá presumir que scatter_indices é classificado (em ordem crescente, depois de espalhar os valores de acordo com scatter_dims_to_operand_dims) pelo usuário. Se não forem, a semântica é definida pela implementação.

Se unique_indices for definido como verdadeiro, a XLA poderá presumir que todos os elementos dispersos são exclusivos. Assim, o XLA pode usar operações não atômicas. Se unique_indices for definido como verdadeiro e os índices para os quais os dados estão sendo espalhados não forem exclusivos, a semântica será definida pela implementação.

Informalmente, a operação de dispersão pode ser considerada um inverso da operação de coleta, ou seja, a operação de dispersão atualiza os elementos na entrada que são extraídos pela operação de coleta correspondente.

Para uma descrição informal detalhada e exemplos, consulte a seção "Descrição informal" em Gather.

Selecionar

Consulte também XlaBuilder::Select.

Constrói uma matriz de saída com base nos valores de duas matrizes de entrada, com base nos valores de uma matriz de predicado.

Select(pred, on_true, on_false)

As matrizes on_true e on_false precisam ter a mesma forma. Essa também é a forma da matriz de saída. A matriz pred precisa ter a mesma dimensionalidade que on_true e on_false, com o tipo de elemento PRED.

Para cada elemento P de pred, o elemento correspondente da matriz de saída é retirado de on_true se o valor de P for true e de on_false se o valor de P for false. Como uma forma restrita de transmissão, pred pode ser um escalar do tipo PRED. Nesse caso, a matriz de saída é tirada inteiramente de on_true se pred for true e de on_false se pred for false.

Exemplo com pred não escalar:

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

Exemplo com pred escalar:

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

As seleções entre tuplas são aceitas. As tuplas são consideradas tipos escalares para essa finalidade. Se on_true e on_false forem tuplas (que precisam ter a mesma forma), pred precisará ser um escalar do tipo PRED.

SelectAndScatter

Consulte também XlaBuilder::SelectAndScatter.

Essa operação pode ser considerada como uma operação composta que primeiro calcula ReduceWindow na matriz operand para selecionar um elemento de cada janela e, em seguida, espalha a matriz source para os índices dos elementos selecionados para construir uma matriz de saída com a mesma forma que a matriz de operandos. A função binária select é usada para selecionar um elemento de cada janela, aplicando-o em cada janela, e é chamada com a propriedade de que o vetor de índice do primeiro parâmetro é lexicograficamente menor que o vetor de índice do segundo parâmetro. A função select retorna true se o primeiro parâmetro for selecionado e retorna false se o segundo parâmetro for selecionado. A função precisa manter a transitividade (ou seja, se select(a, b) e select(b, c) forem true, select(a, c) também será true) para que o elemento selecionado não dependa da ordem dos elementos percorridos para uma determinada janela.

A função scatter é aplicada em cada índice selecionado na matriz de saída. Ela usa dois parâmetros escalares:

1. Valor atual no índice selecionado na matriz de saída
2. O valor de dispersão de source que se aplica ao índice selecionado

Ele combina os dois parâmetros e retorna um valor escalar que é usado para atualizar o valor no índice selecionado na matriz de saída. Inicialmente, todos os índices da matriz de saída são definidos como init_value.

A matriz de saída tem a mesma forma que a matriz operand, e a matriz source precisa ter a mesma forma que o resultado da aplicação de uma operação ReduceWindow na matriz operand. SelectAndScatter pode ser usado para propagar de volta os valores de gradiente de uma camada de agregação em uma rede neural.

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

A figura abaixo mostra exemplos de uso de SelectAndScatter, com a função select calculando o valor máximo entre os parâmetros. Quando as janelas se sobrepõem, como na figura (2) abaixo, um índice da matriz operand pode ser selecionado várias vezes por janelas diferentes. Na figura, o elemento de valor 9 é selecionado pelas duas janelas de cima (azul e vermelho), e a função binária de adição scatter produz o elemento de saída de valor 8 (2 + 6).

A ordem de avaliação da função scatter é arbitrária e pode ser não determinística. Portanto, a função scatter não deve ser muito sensível à rea associação. Consulte a discussão sobre associatividade no contexto de Reduce para mais detalhes.

Enviar

Consulte também XlaBuilder::Send.

Send(operand, channel_handle)

Envia os dados do operando fornecidos para uma instrução Recv em outra computação que compartilha o mesmo identificador de canal. Não retorna dados.

Assim como a operação Recv, a API cliente da operação Send representa a comunicação síncrona e é decomposta internamente em duas instruções HLO (Send e SendDone) para permitir transferências de dados assíncronas. Consulte também HloInstruction::CreateSend e HloInstruction::CreateSendDone.

Send(HloInstruction operand, int64 channel_id)

Inicia uma transferência assíncrona do operando para os recursos alocados pela instrução Recv com o mesmo ID de canal. Retorna um contexto, que é usado por uma instrução SendDone para aguardar a conclusão da transferência de dados. O contexto é um tupla de {operand (shape), request identifier (U32)} e só pode ser usado por uma instrução SendDone.

SendDone(HloInstruction context)

Dado um contexto criado por uma instrução Send, aguarda a conclusão da transferência de dados. A instrução não retorna dados.

Programação de instruções do canal

A ordem de execução das quatro instruções para cada canal (Recv, RecvDone, Send, SendDone) é a seguinte.

• Recv acontece antes de Send
• Send acontece antes de RecvDone
• Recv acontece antes de RecvDone
• Send acontece antes de SendDone

Quando os compiladores de back-end geram uma programação linear para cada cálculo que se comunica por instruções de canal, não pode haver ciclos nas cálculos. Por exemplo, as programações abaixo levam a deadlocks.

A restrição nas instruções se aplica apenas a TPUs no momento da execução. Na GPU, send e recv bloqueiam e não enviam dados reais até que uma solução entre os dispositivos de origem e de destino seja feita.

Fração

Consulte também XlaBuilder::Slice.

O corte extrai um subconjunto da matriz de entrada. A submatriz tem o mesmo número de dimensões que a entrada e contém os valores dentro de uma caixa delimitadora dentro da matriz de entrada, em que as dimensões e índices da caixa delimitadora são fornecidos como argumentos para a operação de fatia.

Slice(operand, start_indices, limit_indices, strides)

Exemplo unidimensional:

let a = {0.0, 1.0, 2.0, 3.0, 4.0}
Slice(a, {2}, {4}) produces:
  {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} }

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

Classificar

Consulte também XlaBuilder::Sort.

Sort(operands, comparator, dimension, is_stable)

Se apenas um operando for fornecido:

• Se o operando for um tensor unidimensional (uma matriz), o resultado será uma matriz ordenada. Se você quiser classificar a matriz em ordem crescente, o comparador precisa realizar uma comparação menor que. Formalmente, depois que a matriz é classificada, ela é mantida para todas as posições de índice i, j com i < j que comparator(value[i], value[j]) = comparator(value[j], value[i]) = false ou comparator(value[i], value[j]) = true.

• Se o operando tiver um número maior de dimensões, ele será classificado de acordo com a dimensão fornecida. Por exemplo, para um tensor bidimensional (uma matriz), um valor de dimensão de 0 classificará cada coluna de forma independente, e um valor de dimensão de 1 classificará cada linha de forma independente. Se nenhum número de dimensão for fornecido, a última dimensão será escolhida por padrão. Para a dimensão ordenada, a mesma ordem de classificação se aplica como no caso unidimensional.

Se os operandos n > 1 forem fornecidos:

• Todos os operandos n precisam ser tensores com as mesmas dimensões. Os tipos de elementos dos tensores podem ser diferentes.

• Todos os operandos são classificados juntos, não individualmente. Conceitualmente, os operandos são tratados como uma tupla. Ao verificar se os elementos de cada operando nas posições de índice i e j precisam ser trocados, o comparador é chamado com parâmetros escalares 2 * n, em que o parâmetro 2 * k corresponde ao valor na posição i do operando k-th, e o parâmetro 2 * k + 1 corresponde ao valor na posição j do operando k-th. Normalmente, o comparador compara os parâmetros 2 * k e 2 * k + 1 entre si e talvez use outros pares de parâmetros como desempates.

• O resultado é uma tupla que consiste nos operandos em ordem classificada (na dimensão fornecida, como acima). O operando i-th da tupla corresponde ao operando i-th da ordenação.

Por exemplo, se houver três operandos operand0 = [3, 1], operand1 = [42, 50] e operand2 = [-3.0, 1.1], e o comparador comparar apenas os valores de operand0 com "menor que", a saída da classificação será a tupla ([1, 3], [50, 42], [1.1, -3.0]).

Se is_stable for definido como verdadeiro, a ordenação será estável. Ou seja, se elementos forem considerados iguais pelo comparador, a ordem relativa dos valores iguais será preservada. Dois elementos e1 e e2 são iguais se e somente se comparator(e1, e2) = comparator(e2, e1) = false. Por padrão, is_stable é definido como falso.

TopK

Consulte também XlaBuilder::TopK.

TopK encontra os valores e índices dos elementos maiores ou menores de k para a última dimensão do tensor fornecido.

TopK(operand, k, largest)

Para um tensor de entrada unidimensional (uma matriz), encontra as entradas k maiores ou menores na matriz e gera uma tupla de duas matrizes (values, indices). Assim, values[j] é a entrada j-ésima maior/menor em operand, e o índice dela é indices[j].

Para um tensor de entrada com mais de uma dimensão, calcula as principais entradas k ao longo da última dimensão, preservando todas as outras dimensões (linhas) na saída. Portanto, para um operando de forma [A, B, ..., P, Q] em que Q >= k, a saída é uma tupla (values, indices) em que:

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

Se dois elementos em uma linha forem iguais, o elemento de índice mais baixo vai aparecer primeiro.

Transposição

Consulte também a operação tf.reshape.

Transpose(operand)

Permuta as dimensões do operando com a permutação fornecida, ou seja, ∀ i . 0 ≤ i < number of dimensions ⇒ input_dimensions[permutation[i]] = output_dimensions[i].

Isso é o mesmo que Reshape(operand, permutation, Permute(permutation, operand.shape.dimensions)).

TriangularSolve

Consulte também XlaBuilder::TriangularSolve.

Resolve sistemas de equações lineares com matrizes de coeficientes triangulares inferiores ou superiores por substituição direta ou inversa. Transmitindo ao longo das dimensões principais, essa rotina resolve um dos sistemas de matriz op(a) * x = b ou x * op(a) = b para a variável x, considerando a e b, em que op(a) é op(a) = a, op(a) = Transpose(a) ou op(a) = Conj(Transpose(a)).

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

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 o número de dimensões de a e b for maior que 2, elas serão tratadas como lotes de matrizes, em que todas, exceto as duas dimensões menores, são dimensões de lote. a e b precisam ter dimensões de lote iguais.

Tupla

Consulte também XlaBuilder::Tuple.

Uma tupla que contém um número variável de identificadores de dados, cada um com sua própria forma.

Isso é análogo a std::tuple 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);

As tuplas podem ser deconstruídas (acessadas) pela operação GetTupleElement.

Embora

Consulte também XlaBuilder::While.

While(condition, body, init)

Executa sequencialmente o body até que o condition falhe. Isso é semelhante a um loop while típico em muitos outros idiomas, exceto pelas diferenças e restrições listadas abaixo.

• Um nó While retorna um valor do tipo T, que é o resultado da última execução do body.
• A forma do tipo T é determinada de forma estática e precisa ser a mesma em todas as iterações.

Os parâmetros T das computações são inicializados com o valor init na primeira iteração e são atualizados automaticamente para o novo resultado de body em cada iteração subsequente.

Um dos principais casos de uso do nó While é implementar a execução repetida de treinamento em redes neurais. O pseudocódigo simplificado é mostrado abaixo com um gráfico que representa a computação. O código pode ser encontrado em while_test.cc. O tipo T neste exemplo é um Tuple que consiste em um int32 para a contagem de iterações e um vector[10] para o acumulador. Para 1.000 iterações, o loop continua adicionando um vetor constante ao acumulador.

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