Semântica de operação

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

Uma observação sobre a nomenclatura: o tipo de dados generalizado que o XLA lida é uma matriz N-dimensional que contém elementos de algum tipo uniforme (por exemplo, flutuante de 32 bits). Em toda a documentação, array é usado para denotar um array 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)

Argumentos	Tipo	Semântica
operands	XlaOp	número variado de tokens

AllGather

Consulte também XlaBuilder::AllGather.

Executa a concatenação entre réplicas.

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

Argumentos	Tipo	Semântica
operand	XlaOp	Matriz a ser concatenada em réplicas
all_gather_dim	int64	Dimensão de concatenação
replica_groups	vetor de vetores de int64	Grupos entre os quais a concatenação é realizada
channel_id	opcional int64	ID do canal opcional para comunicação entre módulos

• 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á vazio.
• O 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)

Argumentos	Tipo	Semântica
operand	XlaOp	Matriz ou uma tupla não vazia de matrizes para reduzir entre as réplicas
computation	XlaComputation	Cálculo da redução
replica_groups	vetor de vetores de int64	Grupos entre os quais as reduções são realizadas
channel_id	opcional int64	ID do canal opcional para comunicação entre módulos

• 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} executa redução entre as réplicas 0 e 2 e 1 e 3.
• O 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 desse cálculo de operação e soma será [4.0, 7.75] em ambas as réplicas. Se a entrada for uma tupla, a saída também será uma tupla.

A computação 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 aguardará 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 split_count de blocos ao longo do 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.

Para configurar os núcleos participantes, faça o seguinte:

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

Argumentos	Tipo	Semântica
operand	XlaOp	Matriz de entrada n-dimensional
split_dimension	int64	Um valor no intervalo [0, n) que nomeia a dimensão em que o operando é dividido
concat_dimension	int64	Um valor no intervalo [0, n) que nomeia a dimensão em que os blocos divididos são concatenados
split_count	int64	O número de núcleos que participam dessa operação. Se replica_groups estiver vazio, esse será o número de réplicas. Caso contrário, ele será igual ao número de réplicas em cada grupo.
replica_groups	Vetor ReplicaGroup	Cada grupo contém uma lista de IDs de réplica.

Confira abaixo um exemplo do 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 0, de modo que 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 ao longo da dimensão 1, 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)

Argumentos	Tipo	Semântica
operand	XlaOp	Matriz n-dimensional a ser normalizada (x)
scale	XlaOp	Matriz unidimensional (\(\gamma\))
mean	XlaOp	Matriz unidimensional (\(\mu\))
variance	XlaOp	Matriz unidimensional (\(\sigma^2\))
grad_output	XlaOp	Gradientes transmitidos para BatchNormTraining (\(\nabla y\))
epsilon	float	Valor de Épsilon (\(\epsilon\))
feature_index	int64	Índice para a dimensão do recurso em operand

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 seguintes fórmulas (supondo 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:

Saídas	Tipo	Semântica
grad_operand	XlaOp	gradiente em relação à entrada operand ($\nabla x$)
grad_scale	XlaOp	gradiente em relação à entrada scale ($\nabla \gamma$)
grad_offset	XlaOp	gradiente em relação à entrada offset($\nabla \beta$)

BatchNormInference

Consulte também XlaBuilder::BatchNormInference e o documento original de normalização em 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)

Argumentos	Tipo	Semântica
operand	XlaOp	matriz n-dimensional a ser normalizada
scale	XlaOp	Matriz unidimensional
offset	XlaOp	Matriz unidimensional
mean	XlaOp	Matriz unidimensional
variance	XlaOp	Matriz unidimensional
epsilon	float	Valor de epsilon
feature_index	int64	Índice para a dimensão do recurso em operand

Para cada recurso na dimensão de recurso (feature_index é o índice da dimensão 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. Em vez disso, 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)

Argumentos	Tipo	Semântica
operand	XlaOp	Matriz n-dimensional a ser normalizada (x)
scale	XlaOp	Matriz unidimensional (\(\gamma\))
offset	XlaOp	Matriz unidimensional (\(\beta\))
epsilon	float	Valor de Épsilon (\(\epsilon\))
feature_index	int64	Índice para a dimensão de recurso em operand

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 funciona da seguinte maneira para cada lote em operand \(x\) que contém elementos m com w e h como o tamanho das dimensões espaciais (supondo que operand seja uma matriz de quatro dimensões):

• Calcula a média do lote \(\mu_l\) para cada atributo l na dimensão: \(\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 desloca: \(y_{ijkl}=\frac{\gamma_l(x_{ijkl}-\mu_l)}{\sqrt[2]{\sigma^2_l+\epsilon} }+\beta_l\)

O valor de épsilon, 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:

Saídas	Tipo	Semântica
output	XlaOp	matriz dimensional com a mesma forma da entrada operand (y)
batch_mean	XlaOp	Matriz unidimensional (\(\mu\))
batch_var	XlaOp	Matriz unidimensional (\(\sigma^2\))

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 da entrada e da saída precisa ser correspondente. Por exemplo, os elementos s32 se tornam elementos f32 pela rotina bitcast, e um elemento s32 se tornará 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)

Argumentos	Tipo	Semântica
operand	XlaOp	matriz do tipo T com escurecimentos D
new_element_type	PrimitiveType	tipo U

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

Argumentos	Tipo	Semântica
operand	XlaOp	A matriz a ser duplicada
broadcast_sizes	ArraySlice<int64>	Os tamanhos das novas dimensões

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

O novo índice de dimensões é indexado 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 a classificação de uma matriz duplicando os dados nela.

BroadcastInDim(operand, out_dim_size, broadcast_dimensions)

Argumentos	Tipo	Semântica
operand	XlaOp	A matriz a ser duplicada
out_dim_size	ArraySlice<int64>	Os tamanhos das dimensões da forma de destino
broadcast_dimensions	ArraySlice<int64>	A dimensão na forma de destino a que cada dimensão da forma do operando corresponde

É 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 no formato de saída para o 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 um cálculo com os argumentos fornecidos.

Call(computation, args...)

Argumentos	Tipo	Semântica
computation	XlaComputation	Cálculo do tipo T_0, T_1, ..., T_{N-1} -> S com N parâmetros de tipo arbitrário
args	sequência de N XlaOps	N argumentos de tipo arbitrário

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

Cholesky

Consulte também XlaBuilder::Cholesky.

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

Cholesky(a, lower)

Argumentos	Tipo	Semântica
a	XlaOp	uma matriz de classificação > 2 de um tipo complexo ou de ponto flutuante.
lower	bool	se o triângulo superior ou inferior de a será usado.

Se lower for true, calcule as matrizes triangulares inferiores l de modo que $a = l . l^T$. Se lower for false, calcule as matrizes triangulares superiores 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 classificação de a for maior que 2, 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étrico (ermitiano) definido como positivo, o resultado será definido para implementação.

Braçadeira

Consulte também XlaBuilder::Clamp.

Limita um operando dentro do intervalo entre um valor mínimo e máximo.

Clamp(min, operand, max)

Argumentos	Tipo	Semântica
min	XlaOp	matriz do tipo T
operand	XlaOp	matriz do tipo T
max	XlaOp	matriz do tipo T

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.

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

Collapse(operand, dimensions)

Argumentos	Tipo	Semântica
operand	XlaOp	matriz do tipo T
dimensions	Vetor int64	subconjunto consecutivo e em ordem das dimensões de T.

A opção de recolher substitui o subconjunto fornecido 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 (a mais principal) no aninhamento de loop que recolhe essas dimensões, e o número de dimensão mais alto tem variação mais rápida (a mais secundária). 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)

Argumentos	Tipo	Semântica
operand	XlaOp	Matriz de entrada n-dimensional
source_target_pairs	Vetor <int64, int64>	Uma lista de pares (source_replica_id, target_replica_id). Para cada par, o operando é enviado da réplica de origem para a réplica de destino.

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 dessa réplica será um tensor que consiste em 0(s) com a mesma forma que a entrada.

Concatenate

Consulte também XlaBuilder::ConcatInDim.

A função "Concatenate" compõe uma matriz usando vários operandos de matriz. A matriz tem a mesma classificação de cada um dos operandos de matriz de entrada (que precisam ter a mesma classificação um do outro) e contém os argumentos na ordem em que foram especificados.

Concatenate(operands..., dimension)

Argumentos	Tipo	Semântica
operands	sequência de N XlaOp	N matrizes do tipo T com dimensões [L0, L1, ...]. Requer N >= 1.
dimension	int64	Um valor no intervalo [0, N) que nomeia a dimensão a ser concatenada entre o operands.

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, os valores de rank 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)

Argumentos	Tipo	Semântica
pred	XlaOp	Escalar do tipo PRED
true_operand	XlaOp	Argumento do tipo \(T_0\)
true_computation	XlaComputation	XlaComputation do tipo \(T_0 \to S\)
false_operand	XlaOp	Argumento do tipo \(T_1\)
false_computation	XlaComputation	XlaComputation do tipo \(T_1 \to S\)

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

A true_computation precisa receber um único argumento do tipo \(T_0\) e será invocada 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)

Argumentos	Tipo	Semântica
branch_index	XlaOp	Escalar do tipo S32
branch_computations	sequência de N XlaComputation	XlaComputations do tipo \(T_0 \to S , T_1 \to S , ..., T_{N-1} \to S\)
branch_operands	sequência de N XlaOp	Argumentos do tipo \(T_0 , T_1 , ..., T_{N-1}\)

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

Argumentos	Tipo	Semântica
lhs	XlaOp	matriz de entradas de rank n+2
rhs	XlaOp	classificar n+2 matriz de pesos de kernel
window_strides	ArraySlice<int64>	Matriz n-d de incrementos de kernel
padding	ArraySlice< pair<int64,int64>>	Matriz n-d de padding (baixo, alto)
lhs_dilation	ArraySlice<int64>	Matriz de fator de dilatação de lhs n-d
rhs_dilation	ArraySlice<int64>	matriz do fator de dilatação n-d rhs
feature_group_count	int64	o número de grupos de atributos
batch_group_count	int64	o número de grupos em lote

Vamos supor que n seja o número de dimensões espaciais. O argumento lhs é uma matriz de classificação n+2 que descreve a área de base. Isso é chamado de entrada, mesmo que o rhs 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 pela qual a janela se move.

O argumento rhs é uma matriz de classificação n+2 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, os buracos d-1 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 dos lhs 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 do 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 são divididas igualmente em vários grupos feature_group_count, que consistem 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 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 recurso de saída.

Para convolução em profundidade, o argumento feature_group_count seria definido como a dimensão do recurso 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 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 em lhs.
• z: 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 de batch, z, y e x na saída. Em seguida, (y,x) é a 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, portanto, 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 pegar a soma dos produtos por elementos entre as duas caixas (semelhante a um produto escalar). 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 que é usada. Há uma caixa 3D separada de valores usada 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)

Argumentos	Tipo	Semântica
operand	XlaOp	matriz do tipo T com escurecimentos D
new_element_type	PrimitiveType	tipo U

As dimensões do operando e o formato 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 realizará uma rotina de conversão int-para-flutuante normalização, como de arredondamento para o ponto uniforme 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 um cálculo.

CustomCall(target_name, args..., shape)

Argumentos	Tipo	Semântica
target_name	string	Nome da função. Uma instrução de chamada será emitida para esse nome de símbolo.
args	sequência de N XlaOps	N argumentos de tipo arbitrário, que serão transmitidos para a função.
shape	Shape	Forma de saída da função

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)

Argumentos	Tipo	Semântica
lhs	XlaOp	matriz do tipo T
rhs	XlaOp	matriz do tipo T

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

Entrada	Saída	Semântica
vetor [n] dot vetor [n]	escalar	produto escalar de vetores
matriz [m x k] dot vetor [k]	vetor [m]	multiplicação de matriz-vetor
matriz [m x k] dot matriz [k x n]	matriz [m x n]	multiplicação matricial

A operação realiza a soma dos produtos na segunda dimensão de lhs (ou a primeira, se tiver classificação 1) e a 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)

Argumentos	Tipo	Semântica
lhs	XlaOp	matriz do tipo T
rhs	XlaOp	matriz do tipo T
dimension_numbers	DotDimensionNumbers	números de dimensão de lote e de contratação

Semelhante ao ponto, mas permite que os números das dimensões de contratação e lote sejam especificados para lhs e rhs.

Campos DotDimensionNumbers	Tipo	Semântica
lhs_contracting_dimensions	int64 repetido	Números de dimensão de contratação de lhs
rhs_contracting_dimensions	int64 repetido	Números de dimensão de contratação de rhs
lhs_batch_dimensions	int64 repetido	lhs números de dimensão em lote
rhs_batch_dimensions	int64 repetido	rhs números de dimensão de lote

A DotGeneral executa a soma dos produtos nas dimensões contratuais 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 contrataçã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} } }

Entrada	Saída	Semântica
[b0, m, k] dot [b0, k, n]	[b0, m, n]	matmul em lote
[b0, b1, m, k] dot [b0, b1, k, n]	[b0, b1, m, n]	matmul em lote

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

DynamicSlice

Consulte também XlaBuilder::DynamicSlice.

O DynamicSlice extrai uma submatriz da matriz de entrada na start_indices dinâmica. O tamanho da fração em cada dimensão é transmitido em size_indices, que especifica o ponto final dos intervalos exclusivos de fração em cada dimensão: [início, início + tamanho). A forma de start_indices precisa ser rank == 1, com o tamanho da dimensão igual à classificação de operand.

DynamicSlice(operand, start_indices, size_indices)

Argumentos	Tipo	Semântica
operand	XlaOp	Matriz n-dimensional do tipo T
start_indices	sequência de N XlaOp	Lista de N números inteiros escalares que contêm os índices iniciais da fatia para cada dimensão. O valor precisa ser maior ou igual a zero.
size_indices	ArraySlice<int64>	Lista de N números inteiros contendo o tamanho da fatia de cada dimensão. Cada valor precisa ser estritamente maior que zero, e o tamanho de início + tamanho precisa ser menor ou igual ao tamanho da dimensão para evitar o tamanho da dimensão do módulo.

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

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

Isso garante que a fração extraída esteja sempre dentro dos limites em relação à matriz de operandos. Se a fração 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. O formato de start_indices precisa ter a classificação == 1, com o tamanho da dimensão igual à classificação de operand.

DynamicUpdateSlice(operand, update, start_indices)

Argumentos	Tipo	Semântica
operand	XlaOp	Matriz n-dimensional do tipo T
update	XlaOp	Matriz N dimensional do tipo T que contém a atualização da fatia. Cada dimensão da forma de atualização precisa ser estritamente maior que zero, e a inicialização + atualização precisa ser menor ou igual ao tamanho do operando de cada dimensão para evitar a geração de índices de atualização fora dos limites.
start_indices	sequência de N XlaOp	Lista de N números inteiros escalares que contêm os índices iniciais da fatia para cada dimensão. O valor precisa ser maior ou igual a zero.

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

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

Isso garante que a fração atualizada esteja sempre dentro dos limites em relação à matriz de operandos. Se a fração 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 destas 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 (O lógico), Xor (XOR 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).

Argumentos	Tipo	Semântica
lhs	XlaOp	operando do lado esquerdo: matriz do tipo T
rhs	XlaOp	operando à direita: matriz do tipo T

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. Nesta variante, as operações entre matrizes de classificações diferentes não são compatíveis, a menos que um dos operandos seja 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 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.

Existe uma variante alternativa com suporte à transmissão de classificação diferente para estas operações:

Op(lhs, rhs, broadcast_dimensions)

Em que Op é o mesmo que acima. Essa variante da operação deve 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 expandir o rank do operando de menor rank até o rank do operando de maior rank. broadcast_dimensions mapeia as dimensões da forma de menor classificação para as dimensões da forma de classificação mais alta. 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 por elemento

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.

Argumentos	Tipo	Semântica
lhs	XlaOp	operando do lado esquerdo: matriz do tipo T
rhs	XlaOp	operando à direita: matriz do tipo T

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 de transmissão de diferentes classificaçõ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 por elemento

O XlaBuilder é compatível com estas funções unárias por elemento:

Abs(operand) Elemento x -> |x| de abs.

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

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

Clz(operand) Contagem por elementos zero à esquerda.

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) Exponencial natural por elemento x -> e^x.

Expm1(operand) Exponenciação natural por elemento menos um x -> e^x - 1.

Floor(operand) Elemento mínimo 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, será retornado 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 para elementos 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) x -> !(x) lógico por elemento.

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

Real(operand) Parte real por elemento 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 valor 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 de 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 por elemento x -> tan(x).

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

Argumentos	Tipo	Semântica
operand	XlaOp	O operando da função

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

Ft

A operação XLA FFT 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.

Consulte também XlaBuilder::Fft.

Argumentos	Tipo	Semântica
operand	XlaOp	A matriz que estamos transformando em Fourier.
fft_type	FftType	Consulte a tabela abaixo.
fft_length	ArraySlice<int64>	As durações no domínio do tempo dos eixos que estão sendo transformados. Isso é necessário principalmente para que o IRFFT dimensione corretamente o eixo mais interno, já que RFFT(fft_length=[16]) tem a mesma forma de saída que RFFT(fft_length=[17]).

FftType	Semântica
FFT	FFT complexo-a-complexo para a frente. A forma não foi alterada.
IFFT	FFT inversa de complexo para complexo. A forma não foi alterada.
RFFT	Encaminhar FFTs reais para complexas. A forma do eixo mais interno é reduzida para fft_length[-1] // 2 + 1 se fft_length[-1] for um valor diferente de zero, omitindo a parte conjugada invertida do sinal transformado além da frequência de Nyquist.
IRFFT	FFT real-para-complexo inverso (ou seja, recebe complexo, retorna real). A forma do eixo mais interno é expandida para fft_length[-1] se fft_length[-1] for um valor diferente de zero, inferindo a parte do sinal transformado além da frequência de Nyquist do conjugado reverso das entradas 1 para fft_length[-1] // 2 + 1.

FFT multidimensional

Quando mais de 1 fft_length é fornecido, isso equivale a aplicar uma cascata de operações 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. Então, 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)

Argumentos	Tipo	Semântica
operand	XlaOp	A matriz da qual estamos coletando.
start_indices	XlaOp	Matriz que contém os índices iniciais das fatias coletadas.
index_vector_dim	int64	A dimensão em start_indices que "contém" os índices iniciais. Confira abaixo uma descrição detalhada.
offset_dims	ArraySlice<int64>	O conjunto de dimensões na forma de saída que são deslocadas em uma matriz separada do operando.
slice_sizes	ArraySlice<int64>	slice_sizes[i] são os limites da fração na dimensão i.
collapsed_slice_dims	ArraySlice<int64>	O conjunto de dimensões em cada fração que são recolhidas. Essas dimensões precisam ter o tamanho 1.
start_index_map	ArraySlice<int64>	Um mapa que descreve como mapear índices em start_indices para índices legais no operando.
indices_are_sorted	bool	Indica se os índices serão classificados pelo autor da chamada.

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 de rank 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, consideramos implicitamente start_indices como uma dimensão 1 à direita. Ou seja, se start_indices tinha a forma [6,7] e index_vector_dim é 2, consideramos implicitamente a forma de start_indices como [6,7,1].

Os limites para a 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 algumas k), vamos escolher os limites de dimensão correspondentes de start_indices.shape, pulando 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 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. Vamos supor que 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 + é uma adição elemento.

remapped_offset_dims é uma função monotônica 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 operandos, 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.

• Diminuímos dinamicamente uma fração com tamanho slice_sizes usando o índice inicial completo.

• Reduzimos a forma da fração ao recolher as dimensões collapsed_slice_dims. Como todas as dimensões de fração recolhidas precisam ter um limite de 1, essa remodelaçã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 usa [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 do lote.

Os índices coletados 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 do lote de saída (dimensões que contêm G0 e G1 no último exemplo) são definidas como as dimensões de saída que não são 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 a classificação 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 rank da matriz de entrada, e um mapeamento explícito determina como o índice precisa ser expandido para ter o mesmo rank 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)

Argumentos	Tipo	Semântica
operand	XlaOp	matriz de entrada dimensional
dimension	int64	Um valor no intervalo [0, n) que especifica a dimensão

SetDimensionSize

Consulte também XlaBuilder::SetDimensionSize.

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

SetDimensionSize(operand, size, dimension)

Argumentos	Tipo	Semântica
operand	XlaOp	matriz de entrada n dimensional.
size	XlaOp	int32, que representa o tamanho dinâmico do ambiente de execução.
dimension	int64	Um valor no intervalo [0, n) que especifica a dimensão.

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)

Argumento	Tipo	Semântica
shape	Shape	Formato dos dados lidos da interface do Infeed. O campo de layout da forma precisa ser definido para corresponder ao layout dos dados enviados ao dispositivo. Caso contrário, o comportamento será indefinido.

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 especificados 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 loopswhile.

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

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

Não há suporte para 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.

Argumentos	Tipo	Semântica
shape	Shape	Forma da matriz criada por Iota()
iota_dimension	int64	A dimensão a ser incrementada.

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)

Argumentos	Tipo	Semântica
operands	sequência de N XlaOps	N matrizes dos tipos T0..T{N-1}
computation	XlaComputation	Cálculo do tipo T_0, T_1, .., T_{N + M -1} -> S com N parâmetros do tipo T e M de tipo arbitrário
dimensions	Matriz int64	matriz de dimensões do mapa

Aplica uma função escalar às matrizes operands fornecidas, 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.

Absorvente

Consulte também XlaBuilder::Pad.

Pad(operand, padding_value, padding_config)

Argumentos	Tipo	Semântica
operand	XlaOp	matriz do tipo T
padding_value	XlaOp	escalar do tipo T para preencher o padding adicionado.
padding_config	PaddingConfig	Valor do padding nas duas bordas (baixa, alta) e entre os elementos de cada dimensão

Expande a matriz operand especificada adicionando padding ao redor dela 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 padding 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 adicionado entre dois elementos em cada dimensão. Ele não pode ser negativo. O padding interno ocorre de maneira lógica antes do padding da borda. Portanto, no caso do padding negativo, os elementos são removidos do operando com preenchimento interno.

Essa operação não é executada 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)

Argumentos	Tipo	Semântica
shape	Shape	a forma dos dados para receber
channel_handle	ChannelHandle	identificador exclusivo para cada par de envio/recebimento

Recebe dados com a 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)

Argumentos	Tipo	Semântica
operands	Sequência de N XlaOp	N matrizes de tipos T_0, ..., T_{N-1}.
init_values	Sequência de N XlaOp	N escalares dos tipos T_0, ..., T_{N-1}.
computation	XlaComputation	computação do tipo T_0, ..., T_{N-1}, T_0, ..., T_{N-1} -> Collate(T_0, ..., T_{N-1}).
dimensions	Matriz int64	matriz desordenada de dimensões a serem reduzidas.

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) será 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. A classificação de cada matriz retornada é rank(operand) - len(dimensions). A saída da operação é Collate(Q_0, ..., Q_N), em que Q_i é uma matriz do tipo T_i, com as dimensões 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 flutuações. No entanto, se o intervalo de dados for limitado, a adição de ponto flutuante será suficientemente próxima para ser associativa para a maioria das utilizações práticas.

Exemplos

Ao reduzir 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 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))))

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 rank of the result
for r0 in range(result_shape[0]), r1 in range(result_shape[1]), ...:
  # Initialize this result element
  result[r0, r1...] <- 0

  # Iterate over all the reduction dimensions
  for d0 in range(dimensions[0]), d1 in range(dimensions[1]), ...:
    # Increment the result element with the value of the operand's element.
    # The index of the operand's element is constructed from all ri's and di's
    # in the right order (by construction ri's and di's together index over the
    # whole operand shape).
    result[r0, r1...] += operand[ri... di]

Este é um exemplo de redução de uma matriz 2D. A forma tem classificação 2, dimensão 0 de tamanho 2 e 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. Para facilitar a visualização, o diagrama mostra uma como coluna e outra como linha.

Para um exemplo mais complexo, confira uma matriz 3D. A classificação é 3, a dimensão 0 tem tamanho 4, a dimensão 1 tem tamanho 2 e a dimensão 2 tem 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 de rank 2 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 teremos uma matriz de classificação 2 em que todos os valores na 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 a classificação muda.

Também podemos 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.

Redução variável

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:

• Execução de valor reduzido para o primeiro operando
• …
• Executando o valor reduzido para o N-ésimo operando
• Valor de entrada para o primeiro operando
• …
• Valor de entrada para o no 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 1-D 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 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 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)

Argumentos	Tipo	Semântica
operand	XlaOp	matriz do tipo de ponto flutuante T.
exponent_bits	int32	número de bits de expoente no formato de precisão inferior
mantissa_bits	int32	número de bits de mantissa em formato de menor precisão

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 determinado de bits de mantissa (usando a semântica "empates para pares"), e todos os valores que excedem o intervalo especificado pelo número de bits expoentes são fixados ao 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 expoente (para distinguir um valor zero de um infinito, já que ambos têm zero mantissa) 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 um ambiente autônomo.

ReduceScatter

Consulte também XlaBuilder::ReduceScatter.

ReduzirScatter é uma operação coletiva que efetivamente faz um AllReduce e espalha 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)

Argumentos	Tipo	Semântica
operand	XlaOp	Matriz ou uma tupla não vazia de matrizes para reduzir entre as réplicas.
computation	XlaComputation	Cálculo da redução
scatter_dimension	int64	Dimensão para dispersão.
shard_count	int64	Número de blocos para dividir scatter_dimension
replica_groups	vetor de vetores de int64	Grupos entre os quais as reduções são realizadas
channel_id	opcional int64	ID do canal opcional para comunicação entre módulos

• Quando operand é uma tupla de matrizes, a redução de dispersão é 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. 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 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á vazio. Se replica_groups não estiver vazio, shard_count precisa 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.

O formato de saída é o de entrada, em que 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 pool pode ser expressa como um ReduceWindow. Semelhante a Reduce, o computation aplicado é sempre transmitido para o init_values no lado esquerdo.

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

Argumentos	Tipo	Semântica
operands	N XlaOps	Uma sequência de N matrizes multidimensionais do tipo T_0,..., T_{N-1}, cada uma representando a área de base em que a janela está posicionada.
init_values	N XlaOps	Os N valores iniciais para a redução, um para cada um dos N operandos. Consulte Reduzir para mais detalhes.
computation	XlaComputation	Função de redução do tipo T_0, ..., T_{N-1}, T_0, ..., T_{N-1} -> Collate(T_0, ..., T_{N-1}), para aplicar aos elementos em cada janela de todos os operandos de entrada.
window_dimensions	ArraySlice<int64>	matriz de números inteiros para valores de dimensão da janela
window_strides	ArraySlice<int64>	matriz de números inteiros para valores de janela de passo
base_dilations	ArraySlice<int64>	matriz de números inteiros para valores de dilatação de base
window_dilations	ArraySlice<int64>	matriz de números inteiros para valores de dilatação da janela
padding	Padding	Tipo de preenchimento para a janela (Padding::kSame, que preenche para ter a mesma forma de saída que a entrada se o passo for 1, ou Padding::kValid, que não usa preenchimento e "interrompe" a janela quando ela não se encaixa mais)

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

Um 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 preenchimento é aplicado a cada dimensão da entrada, e os cálculos são os mesmos que se a entrada tivesse vindo com as dimensões que ela tem após o preenchimento.

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 sobre o array 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 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, new_sizes) Reshape(operand, dimensions, new_sizes)

Argumentos	Tipo	Semântica
operand	XlaOp	matriz do tipo T
dimensions	Vetor int64	ordem em que as dimensões são recolhidas
new_sizes	Vetor int64	vetor de tamanhos de novas dimensões

Conceitualmente, a remodelação primeiro nivela 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. Os valores no vetor dimension, se fornecidos, precisam ser uma permutação de todas as dimensões de T. O padrão, se não for fornecido, é {0, ..., rank - 1}. A ordem das dimensões em dimensions vai da dimensão de variação mais lenta (mais importante) à dimensão de variação mais rápida (mais secundária) no loop aninhado, que reduz a matriz de entrada a uma única dimensão. O vetor new_sizes determina o tamanho da matriz de saída. O valor no índice 0 em new_sizes é 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 new_size precisa ser igual ao produto dos tamanhos de dimensão do operando. Ao refinar a matriz recolhida na matriz multidimensional definida por new_sizes, as dimensões em new_sizes são ordenadas da variação mais lenta (a mais principal) e à variável mais rápida (a 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} } };

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

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

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

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


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

Como 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} }, {0,1}, {}) == 5;
Reshape(5, {}, {1,1}) == f32[1x1] { {5} };

Rev (reverso)

Consulte também XlaBuilder::Rev.

Rev(operand, dimensions)

Argumentos	Tipo	Semântica
operand	XlaOp	matriz do tipo T
dimensions	ArraySlice<int64>	dimensões a serem invertidas

Reverte a ordem dos elementos na matriz operand ao longo do dimensions especificado, gerando uma matriz de saída da 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 a forma de saída precisam ter um tipo elementar de ponto flutuante. Além disso, os parâmetros precisam ter valores escalares.

RngNormal(mu, sigma, shape)

Argumentos	Tipo	Semântica
mu	XlaOp	Escalar do tipo T especificando a média dos números gerados
sigma	XlaOp	Escalar do tipo T especificando o desvio padrão dos dados gerados
shape	Shape	Forma de saída do tipo T

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. Atualmente, 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 valores escalares. Se \(b <= a\) , o resultado é definido pela implementação.

RngUniform(a, b, shape)

Argumentos	Tipo	Semântica
a	XlaOp	Escalar do tipo T que especifica o limite inferior do intervalo
b	XlaOp	Escalar do tipo T que especifica o limite superior do intervalo
shape	Shape	Forma de saída do tipo T

RngBitGenerator

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

O estado inicial é o estado inicial da geração de número aleatório 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)

Argumentos	Tipo	Semântica
algorithm	RandomAlgorithm	Algoritmo PRNG a ser usado.
initial_state	XlaOp	Estado inicial do algoritmo PRNG.
shape	Shape	Forma de saída para dados gerados.

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: tão fáceis quanto 1, 2, 3.

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

Argumentos	Tipo	Semântica
operands	Sequência de N XlaOp	N matrizes de tipos T_0, ..., T_N a serem espalhadas.
scatter_indices	XlaOp	Matriz contendo os índices iniciais das fatias que precisam ser dispersas.
updates	Sequência de N XlaOp	N matrizes dos tipos T_0, ..., T_N. updates[i] contém os valores que precisam ser usados para distribuir operands[i].
update_computation	XlaComputation	Cálculo a ser usado para combinar os valores existentes na matriz de entrada e as atualizações durante a dispersão. Essa computação precisa ser do tipo T_0, ..., T_N, T_0, ..., T_N -> Collate(T_0, ..., T_N).
index_vector_dim	int64	A dimensão em scatter_indices que contém os índices iniciais.
update_window_dims	ArraySlice<int64>	O conjunto de dimensões no formato updates que são dimensões de janela.
inserted_window_dims	ArraySlice<int64>	O conjunto de dimensões da janela que precisa ser inserido na forma updates.
scatter_dims_to_operand_dims	ArraySlice<int64>	Uma dimensão é mapeada dos índices de dispersão para o espaço do índice do operando. Essa matriz é interpretada como um mapeamento de i para scatter_dims_to_operand_dims[i] . Tem que ser individual e total.
indices_are_sorted	bool	Se os índices são garantidos para serem classificados pelo autor da chamada.
unique_indices	bool	Se os índices são garantidos como exclusivos pelo autor da chamada.

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) será 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 o nível 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 contabilizar inserted_window_dims (ou seja, adjusted_window_bounds[k], em que adjusted_window_bounds contém os limites de operand com os limites em í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 ter números de dimensão repetidos e estar no intervalo [0, updates.rank).

• inserted_window_dims precisa estar em ordem crescente, não ter números de dimensão repetidos e 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 correspondente I na matriz operands correspondente em que essa atualização precisa ser aplicada é calculado da seguinte maneira:

1. Seja 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 ao distribuir S pelo 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 é determinista.

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

Se indices_are_sorted for definido como verdadeiro, a XLA poderá presumir que scatter_indices forem classificados (em ordem crescente, depois de espalhar os valores de acordo com scatter_dims_to_operand_dims) pelo usuário. Caso contrário, a semântica será 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 conferir 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 a partir de elementos de duas matrizes de entrada, com base nos valores de uma matriz de predicados.

Select(pred, on_true, on_false)

Argumentos	Tipo	Semântica
pred	XlaOp	matriz do tipo PRED
on_true	XlaOp	matriz do tipo T
on_false	XlaOp	matriz do tipo T

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 (link em inglês), 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};

É possível fazer seleções entre tuplas. As tuplas são consideradas tipos escalares para esse propósito. Se on_true e on_false forem tuplas (que precisam ter a mesma forma), pred precisa 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 operando. 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. Assim, o elemento selecionado não depende 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 do í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 retropropagar 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)

Argumentos	Tipo	Semântica
operand	XlaOp	matriz do tipo T sobre a qual as janelas deslizam
select	XlaComputation	A computação binária do tipo T, T -> PRED, para ser aplicada a todos os elementos em cada janela. Ela retorna true se o primeiro parâmetro for selecionado e false se o segundo for selecionado.
window_dimensions	ArraySlice<int64>	matriz de números inteiros para valores de dimensão da janela
window_strides	ArraySlice<int64>	matriz de números inteiros para valores de janela de passo
padding	Padding	Tipo de padding para a janela (Padding::kSame ou Padding::kValid)
source	XlaOp	matriz do tipo T com os valores para dispersão
init_value	XlaOp	valor escalar do tipo T para o valor inicial da matriz de saída
scatter	XlaComputation	computação binária do tipo T, T -> T, para aplicar cada elemento de origem de dispersão ao elemento de destino

A figura abaixo mostra exemplos de como usar SelectAndScatter, com a função select calculando o valor máximo entre os parâmetros. Observe que, 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 superiores (azul e vermelha), e a função de adição binária scatter produz o elemento de saída do 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 pode ser muito sensível à reassociação. Consulte a discussão sobre associatividade no contexto de Reduce para mais detalhes.

Enviar

Consulte também XlaBuilder::Send.

Send(operand, channel_handle)

Argumentos	Tipo	Semântica
operand	XlaOp	dados a serem enviados (matriz do tipo T)
channel_handle	ChannelHandle	identificador exclusivo para cada par de envio/recebimento

Envia os dados do operando especificados 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 a seguir para aguardar a conclusão da transferência de dados. O contexto é uma tupla de {operando (forma), identificador de solicitação (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 nos cálculos. Por exemplo, as programações abaixo levam a impasses.

Fração

Consulte também XlaBuilder::Slice.

O fracionamento extrai uma submatriz da matriz de entrada. A submatriz tem a mesma classificação que a entrada e contém os valores dentro de uma caixa delimitadora dentro da matriz de entrada, em que as dimensões e os índices da caixa delimitadora são fornecidos como argumentos para a operação de fração.

Slice(operand, start_indices, limit_indices, strides)

Argumentos	Tipo	Semântica
operand	XlaOp	Matriz n-dimensional do tipo T
start_indices	ArraySlice<int64>	Lista de N números inteiros que contém os índices iniciais da fatia para cada dimensão. Os valores precisam ser maiores ou iguais a zero.
limit_indices	ArraySlice<int64>	Lista de N números inteiros que contêm os índices finais (exclusivos) da fatia para cada dimensão. Cada valor precisa ser maior ou igual ao valor de start_indices correspondente da dimensão e menor ou igual ao tamanho dela.
strides	ArraySlice<int64>	Lista de N números inteiros que decide o passo de entrada da fatia. A fatia seleciona todos os elementos strides[d] na dimensão d.

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)

Argumentos	Tipo	Semântica
operands	ArraySlice<XlaOp>	Os operandos a serem classificados.
comparator	XlaComputation	A computação do comparador a ser usada.
dimension	int64	A dimensão pela qual classificar.
is_stable	bool	Indica se a classificação estável precisa ser usada.

Se apenas um operando for fornecido:

• Se o operando for um tensor de classificação 1 (uma matriz), o resultado será uma matriz classificada. Se você quiser classificar a matriz em ordem crescente, o comparador precisará fazer uma comparação menor que. Formalmente, depois que a matriz é classificada, ela mantém 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 uma classificação mais alta, ele será classificado de acordo com a dimensão fornecida. Por exemplo, para um tensor de ordem 2 (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 classificada, a mesma ordem de classificação é aplicada como no caso de classificação 1.

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

Transposição

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

Transpose(operand)

Argumentos	Tipo	Semântica
operand	XlaOp	O operando a ser transposto.
permutation	ArraySlice<int64>	Como permutar as dimensões.

Permuta as dimensões do operando com a permutação especificada, portanto, ∀ i . 0 ≤ i < rank ⇒ 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 traseira. 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)

Argumentos	Tipo	Semântica
a	XlaOp	uma matriz de classificação > 2 de um tipo complexo ou de ponto flutuante com a forma [..., M, M].
b	XlaOp	uma matriz de rank > 2 do mesmo tipo com o formato [..., M, K] se left_side for verdadeiro, [..., K, M], caso contrário.
left_side	bool	indica se um sistema no formato op(a) * x = b (true) ou x * op(a) = b (false) precisa ser resolvido.
lower	bool	se quer usar o triângulo superior ou inferior de a.
unit_diagonal	bool	Se for true, os elementos diagonais de a serão considerados 1 e não serão acessados.
transpose_a	Transpose	se a será usado como está, transposto ou com a transposta conjugada.

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 classificação 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)

Argumentos	Tipo	Semântica
condition	XlaComputation	XlaComputation do tipo T -> PRED, que define a condição de término do loop.
body	XlaComputation	XlaComputation do tipo T -> T, que define o corpo do loop.
init	T	Valor inicial do parâmetro de condition e body.

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 dos cálculos são inicializados com o valor init na primeira iteração e 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};
}