Veja a seguir a semântica das operações definidas na interface
XlaBuilder
. Normalmente, essas operações são mapeadas individualmente para operações definidas na
interface de 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, como um ponto flutuante de 32 bits. Em toda a documentação, array é usada para indicar uma matriz de dimensão arbitrária. Por 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 usa um número variado de tokens e produz um único token. Os tokens
são tipos primitivos que podem ser encadeados entre operações de efeito colateral para
aplicar a ordenação. AfterAll
pode ser usado como uma junção de tokens para ordenar uma
operação após um conjunto de operações.
AfterAll(operands)
Argumentos | Tipo | Semântica |
---|---|---|
operands |
XlaOp |
número variado de tokens |
AllGather
Consulte também
XlaBuilder::AllGather
.
Executa a concatenação entre as réplicas.
AllGather(operand, all_gather_dim, shard_count, replica_group_ids,
channel_id)
Argumentos | Tipo | Semântica |
---|---|---|
operand
|
XlaOp
|
Matriz para concatenar entre as 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 é executada |
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 é executada (o ID da réplica atual pode ser recuperado usandoReplicaId
). A ordem das réplicas em cada grupo determina a ordem em que as entradas estão localizadas no resultado.replica_groups
precisa estar vazio (nesse caso, todas as réplicas pertencem a um único grupo, ordenadas de0
aN - 1
) ou conter o mesmo número de elementos que o número de réplicas. Por exemplo,replica_groups = {0, 2}, {1, 3}
executa concatenação entre as réplicas0
e2
e1
e3
.shard_count
é o tamanho de cada grupo de réplicas. Precisamos disso nos casos em quereplica_groups
está vazio.- O
channel_id
é usado para comunicação entre módulos: apenas operaçõesall-gather
com o mesmochannel_id
podem se comunicar entre si.
A forma de saída é a forma de entrada com o all_gather_dim
transformado em shard_count
vezes maior. Por exemplo, se houver duas réplicas e o operando tiver o
valor [1.0, 2.5]
e [3.0, 5.25]
, respectivamente, nas duas réplicas, o
valor de saída dessa operação em que all_gather_dim
é 0
será [1.0, 2.5, 3.0,
5.25]
nas duas réplicas.
AllReduce
Consulte também
XlaBuilder::AllReduce
.
Executa um cálculo personalizado entre 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 réplicas |
computation |
XlaComputation |
Computação de 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 atual pode ser recuperado usandoReplicaId
).replica_groups
precisa estar vazio (nesse caso, todas as réplicas pertencem a um único grupo) ou conter o mesmo número de elementos que o número de réplicas. Por exemplo,replica_groups = {0, 2}, {1, 3}
executa a redução entre as réplicas0
e2
e1
e3
.- O
channel_id
é usado para comunicação entre módulos: apenas operaçõesall-reduce
com o mesmochannel_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á.
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 esperará para sempre. Como as réplicas estão executando o mesmo
programa, não há muitas maneiras de isso acontecer, mas é possível quando
a condição de um loop depende dos dados da entrada e os dados que são alimentados
fazem com que a repetição "while" seja iterada 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. Ele tem duas fases:
- A fase de dispersão. Em cada núcleo, o operando é dividido em um número
split_count
de blocos ao longo dasplit_dimensions
e os blocos são espalhados para todos os núcleos. Por exemplo, o i-ésimo bloco é enviado para o i-ésimo núcleo. - A fase de coleta. Cada núcleo concatena os blocos recebidos ao longo do
concat_dimension
.
Os núcleos participantes podem ser configurados da seguinte maneira:
replica_groups
: cada ReplicaGroup contém uma lista de IDs de réplica que participam do cálculo. É possível recuperar o ID da réplica atual usandoReplicaId
. AllToAll será aplicado em 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, os blocos recebidos serão concatenados na mesma ordem de 1, 2, 3. Em seguida, outro AllToAll será aplicado nas réplicas 4, 5, 0, e a ordem de concatenação também será 4, 5, 0. Sereplica_groups
estiver vazio, todas as réplicas pertencerão a um grupo, na ordem de concatenação em que aparecem.
Pré-requisitos:
- O tamanho da dimensão do operando em
split_dimension
é divisível porsplit_count
. - O formato 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
ao lado da qual o operando é
dividido |
concat_dimension
|
int64
|
Um valor no intervalo [0,
n) que nomeia a dimensão
com a qual os blocos de divisão
são concatenados |
split_count
|
int64
|
O número de núcleos que participam dessa operação. Se
replica_groups estiver vazio, esse
precisará ser o número de
réplicas. Caso contrário, deverá 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 na dimensão 0, de modo que cada parte tem a forma f32[4,4]. Elas estão dispersas em todos os núcleos. Em seguida, cada núcleo concatena as partes recebidas na dimensão 1, na ordem do núcleo 0 a 4. Assim, 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 em lote para uma descrição detalhada do algoritmo.
Calcula gradientes da 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 do épsilon (\(\epsilon\)) |
feature_index |
int64 |
Índice para dimensão de atributo em operand |
Para cada recurso na dimensão (feature_index
é o índice da
dimensão 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 do 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 espaciais.
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 de lote e espaciais.
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 épsilon |
feature_index |
int64 |
Índice para dimensão de atributo em operand |
Para cada atributo na dimensão (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 do
recurso em operand
.
BatchNormInference
equivale a chamar BatchNormTraining
sem calcular mean
e variance
para cada lote. Ele usa as entradas mean
e
variance
como valores estimados. O objetivo dessa operação é reduzir a latência na inferência, por isso o nome BatchNormInference
.
A saída é uma matriz normalizada de dimensão n com a mesma forma que a operand
de entrada.
BatchNormTraining
Consulte também
XlaBuilder::BatchNormTraining
e the original batch normalization paper
para ver uma descrição detalhada do algoritmo.
Normaliza uma matriz em dimensões de lote e espaciais.
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 do épsilon (\(\epsilon\)) |
feature_index |
int64 |
Índice para dimensão de atributo em operand |
Para cada atributo na dimensão (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 do
recurso 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 do atributo: \(\mu_l=\frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h x_{ijkl}\)Calcula a variação 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 é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 XlaOp
s:
Saídas | Tipo | Semântica |
---|---|---|
output
|
XlaOp
|
matriz n dimensional com a mesma forma que a entrada operand (y) |
batch_mean |
XlaOp |
Matriz unidimensional (\(\mu\)) |
batch_var |
XlaOp |
Matriz unidimensional (\(\sigma^2\)) |
batch_mean
e batch_var
são momentos calculados nas dimensões de lote e
espaciais usando as fórmulas acima.
BitcastConvertType
Consulte também
XlaBuilder::BitcastConvertType
.
De forma semelhante a um tf.bitcast
no TensorFlow, executa uma operação de bitcast
com elementos de uma forma de dados para uma forma de destino. O tamanho da entrada e da saída precisam
ser correspondentes. Por exemplo, os elementos s32
se tornam elementos f32
pela rotina de bitcast, e um
elemento s32
vai se tornar quatro elementos s8
. O Bitcast é implementado como um cast de baixo nível. Portanto, máquinas com diferentes representações de ponto flutuante fornecerão resultados diferentes.
BitcastConvertType(operand, new_element_type)
Argumentos | Tipo | Semântica |
---|---|---|
operand |
XlaOp |
matriz do tipo T com escurece D |
new_element_type |
PrimitiveType |
tipo U |
As dimensões do operando e o formato de destino precisam ser correspondentes, com exceção da última dimensão, que será alterada pela proporção do tamanho primitivo antes e depois da conversão.
Os tipos de elemento de origem e de destino não podem ser tuplas.
Conversão de Bitcast em um tipo primitivo de largura diferente
A instrução BitcastConvert
do HLO 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 a
operação inteira é conceitualmente um bitcast e não altera os bytes
subjacentes, a forma do elemento de saída precisa mudar. Para B = sizeof(T), B' =
sizeof(T')
, há dois casos possíveis.
Primeiro, quando B > B'
, o formato de saída recebe uma nova dimensão secundária de tamanho
B/B'
. Exemplo:
f16[10,2]{1,0} %output = f16[10,2]{1,0} bitcast-convert(f32[10]{0} %input)
A regra permanece a mesma para escalares efetivos:
f16[2]{0} %output = f16[2]{0} bitcast-convert(f32[] %input)
Como alternativa, para B' > B
, a instrução exige que a última dimensão lógica
do formato de entrada seja igual a B'/B
, e essa dimensão será descartada durante
a conversão:
f32[10]{0} %output = f32[10]{0} bitcast-convert(f16[10,2]{1,0} %input)
As conversões entre diferentes larguras de bits não são compostas por elementos.
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}
.
As novas dimensões são indexadas em cópias do operando, ou seja,
output[i0, ..., iN, j0, ..., jM] = operand[j0, ..., jM]
Por exemplo, se operand
for um f32
escalar com o valor 2.0f
e broadcast_sizes
for {2, 3}
, o resultado será uma matriz com a forma f32[2, 3]
, e todos os valores no resultado serão 2.0f
.
BroadcastInDim
Consulte também
XlaBuilder::BroadcastInDim
.
Expande o tamanho e 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 do formato de destino |
broadcast_dimensions |
ArraySlice<int64> |
A que dimensão na forma de destino cada dimensão do formato do operando corresponde |
Semelhante à transmissão, mas permite adicionar dimensões em qualquer lugar e expandir dimensões existentes com tamanho 1.
O operand
é transmitido para a forma descrita por out_dim_size
.
broadcast_dimensions
mapeia as dimensões de operand
para as dimensões do
formato de destino, ou seja, a i-ésima dimensão do operando é mapeada para a
transmissão_dimension[i]a dimensão da forma de saída. As dimensões de
operand
precisam ter tamanho 1 ou ser do mesmo tamanho que a dimensão na forma
de saída em que são mapeadas. As dimensões restantes são preenchidas com as dimensões de tamanho 1. A transmissão de degeneração de dimensão, em seguida, transmite essas dimensões degeneradas
para alcançar a forma de saída. A semântica é descrita em detalhes na
página de transmissão.
Call
Consulte também
XlaBuilder::Call
.
Invoca um cálculo com os argumentos fornecidos.
Call(computation, args...)
Argumentos | Tipo | Semântica |
---|---|---|
computation |
XlaComputation |
computação do tipo T_0, T_1, ..., T_{N-1} -> S com N parâmetros de tipo arbitrário |
args |
sequência de N XlaOp s |
N argumentos de tipo arbitrário |
A aridade e os tipos do args
precisam corresponder aos parâmetros de
computation
. É permitido não ter args
.
Cholesky
Consulte também
XlaBuilder::Cholesky
.
Calcula a decomposição do 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
, será calculada as matrizes triangulares inferiores l
de modo que $a = l .
l^T$. Se lower
for false
, calculará 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, e 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 as dimensões, exceto as 2 secundárias, são dimensões em lote.
Se a
não for definido como positivo simétrico (Hermitiano), o resultado será
definido pela implementação.
Pinça
Consulte também
XlaBuilder::Clamp
.
Fixa 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 |
Considerando um operando e valores mínimos e máximos, retorna o operando se 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 alternativa, como uma forma restrita de
transmissão, min
e/ou max
podem ser um escalar do tipo T
.
Exemplo com min
e max
escalares:
let operand: s32[3] = {-1, 5, 9};
let min: s32 = 0;
let max: s32 = 6;
==>
Clamp(min, operand, max) = s32[3]{0, 5, 6};
Fechar
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. |
O recolhimento substitui o subconjunto especificado das dimensões do operando por uma única
dimensão. Os argumentos de entrada são uma matriz arbitrária do tipo T e um vetor de constante de compilação e tempo de compilação de índices de dimensão. Os índices de dimensão precisam ser
um subconjunto consecutivo das dimensões de
T em ordem (números de dimensão baixa a alta). Assim, {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 da dimensão, em que o novo
tamanho é igual ao produto dos tamanhos originais. O número de dimensão mais baixo em dimensions
é a dimensão variável mais lenta (a mais importante) no aninhamento de loop, que recolhe essas dimensões. Já o número de dimensão mais alto varia mais rápido (mais menor). Consulte o operador tf.reshape
se uma ordem de recolhimento mais geral for necessária.
Por exemplo, deixe v ser 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
.
CollectivePermute é uma operação coletiva que envia e recebe réplicas de dados entre diferentes.
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 de destino. |
Há as seguintes restrições no 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.
- Caso um ID de réplica não seja um destino em nenhum par, a saída nessa réplica será um tensor que consiste em 0(s) com a mesma forma que a entrada.
Concatenate
Consulte também
XlaBuilder::ConcatInDim
.
"Concatenate" compõe uma matriz com base em vários operandos de matriz. A matriz tem a mesma classificação que cada um dos operandos da matriz de entrada (que precisam ter a mesma classificação entre si) 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 o XLA não oferece suporte a matrizes "ragged". Além disso, os valores de classificação 0 não podem ser concatenados (já que é 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
for false
e retorna o resultado.
O true_computation
precisa ter um único argumento do tipo \(T_0\) e vai ser
invocado com true_operand
, que precisa ser do mesmo tipo. O
false_computation
precisa ter 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 entre 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
< 0 ou >= N, branch_computations[N-1]
será executado como a ramificação padrão.
Cada branch_computations[b]
precisa ter 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.
Observe que 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 VÁLIDO. O MESMO padding preenche a entrada (lhs
) com zeros para que
a saída tenha a mesma forma que a entrada quando não for considerado
o avanço. Padding VÁLIDO significa 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 da base n-dimensional. É realizado um cálculo para cada posição possível da janela.
Argumentos | Tipo | Semântica |
---|---|---|
lhs |
XlaOp |
matriz de entradas de classificação n+2 |
rhs |
XlaOp |
Classificar a matriz n+2 de pesos de kernel |
window_strides |
ArraySlice<int64> |
Matriz n-d de passos do kernel |
padding |
ArraySlice< pair<int64,int64>> |
Matriz n-d de padding (baixo, alto) |
lhs_dilation |
ArraySlice<int64> |
Matriz de fatores de dilatação n-d lhs |
rhs_dilation |
ArraySlice<int64> |
Matriz de fatores 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 |
Permita 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, embora,
claro, os rhs também sejam 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 espaciaisn
que definem a área de base em que a janela se move.
O argumento rhs
é uma matriz de classificação n+2 que descreve o filtro convolucional/kernel/janela. As dimensões estão, nesta ordem:
output-z
: a dimensãoz
da saída.input-z
: o tamanho dessa dimensão vezesfeature_group_count
precisa ser igual ao tamanho da dimensãoz
em lhs.spatial_dims
: descreve as dimensões espaciaisn
que definem a janela "n-d" que se move pela área de base.
O argumento window_strides
especifica o salto da janela convolucional nas dimensões espaciais. Por exemplo, se o salto 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 aplicado à á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 da dimensão y
, e padding[1]
especifica o padding da 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 dos índices mais baixos, enquanto o
alto é aplicado na direção dos índices mais altos. Por exemplo, se padding[1]
for (2,3)
, haverá um preenchimento por dois zeros à esquerda e por três zeros à direita na segunda dimensão espacial. O uso de 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 a lhs e rhs, 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 de ambiente autônomo, que significa zeros para convolução.
A dilatação das rhs também é chamada de convolução atrosa. Para mais detalhes, consulte
tf.nn.atrous_conv2d
. A dilatação do LLM também é chamada de
convolução transposta. Confira mais detalhes em tf.nn.conv2d_transpose
.
O argumento feature_group_count
(valor padrão 1) pode ser usado para convoluções agrupadas. feature_group_count
precisa ser um divisor da dimensão de atributo de entrada e
de saída. Se feature_group_count
for maior que 1, 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 uniformemente em vários grupos de feature_group_count
, cada um consistindo em uma subsequência consecutiva de recursos. A
dimensão de recurso de entrada de rhs
precisa ser igual à dimensão do recurso de entrada
lhs
dividida por feature_group_count
. Portanto, ela já tem o tamanho de um
grupo de atributos de entrada. Os grupos i-th 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 atributo de saída.
Para a convolução de 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 do 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.
O formato de saída tem estas dimensões, nesta ordem:
batch
: o tamanho dessa dimensão vezesbatch_group_count
precisa ser igual ao tamanho da dimensãobatch
em lhs.z
: mesmo tamanho queoutput-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. Efetivamente, dividimos cada lote 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 entre pares e concatenamos a saída com a dimensão do atributo de saída. A semântica operacional de todas as outras dimensões (atributo e espacial) permanece a mesma.
Os posicionamentos válidos da janela convolucional são determinados 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. (y,x)
é uma posição de um canto da janela dentro da área de base (por exemplo, o canto superior esquerdo, dependendo de como você interpreta as dimensões espaciais). Agora temos uma janela 2D, tirada da área de base, em que cada ponto 2D está associado a um vetor 1D, então temos uma caixa 3D. No kernel convolucional, 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 tomar a soma dos produtos por elemento entre as duas
caixas (semelhante a um produto escalar). Esse é o valor de saída.
Observe que, se output-z
for, por exemplo, 5, cada posição da janela produz cinco
valores para a dimensão z
da saída. Esses valores diferem em qual parte do kernel convolucional é usada. Há uma caixa 3D separada de valores usados para cada coordenada output-z
. Você pode pensar nisso como 5 convoluções
separadas com um filtro diferente para cada uma.
Este é um pseudocódigo para uma convolução 2D com preenchimento e salto:
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
.
De forma semelhante a uma static_cast
com elementos em C++, executa uma operação de conversão
de elemento de um formato de dados para um de destino. As dimensões precisam ser correspondentes, e a conversão é por elemento. Por exemplo, os elementos s32
se tornam elementos f32
por meio de uma rotina de conversão de s32
para f32
.
ConvertElementType(operand, new_element_type)
Argumentos | Tipo | Semântica |
---|---|---|
operand |
XlaOp |
matriz do tipo T com escurece D |
new_element_type |
PrimitiveType |
tipo U |
As dimensões do operando e o formato de destino precisam ser correspondentes. Os tipos de elemento de origem e destino não podem ser tuplas.
Uma conversão de T=s32
para U=f32
vai realizar uma rotina de normalização de conversão int para flutuante, como arredondamento para par mais próximo.
let a: s32[3] = {0, 1, 2};
let b: f32[3] = convert(a, f32);
then b == f32[3]{0.0, 1.0, 2.0}
CrossReplicaSum
Executa AllReduce
com um cálculo de soma.
CustomCall
Consulte também
XlaBuilder::CustomCall
.
Chame 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. Será emitida uma instrução de chamada direcionada a esse nome de símbolo. |
args |
sequência de N XlaOp s |
N argumentos de tipo arbitrário, que serão passados para a função. |
shape |
Shape |
Forma de saída da função |
A assinatura da função é a mesma, independentemente da aridade ou do tipo de argumentos:
extern "C" void target_name(void* out, void** in);
Por exemplo, se CustomCall for usado 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 vetor |
matriz [m x k] dot vetor [k] |
vetor [m] | multiplicação de matrizes-vetor |
matriz [m x k] dot matriz [k x n] |
matriz [m x n] | multiplicação de matrizes |
A operação executa a soma de 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 "contratadas". As dimensões contratadas de lhs
e rhs
precisam ser do mesmo tamanho. Na prática, ela pode ser usada para executar 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 |
contratantes e números de dimensão de lote |
Semelhante a ponto, mas permite que os números de dimensão de contrato e lote sejam
especificados para lhs
e rhs
.
Campos DotDimensionNumbers | Tipo | Semântica |
---|---|---|
lhs_contracting_dimensions
|
int64 repetido | lhs números de dimensão de
contração |
rhs_contracting_dimensions
|
int64 repetido | rhs números de dimensão de
contração |
lhs_batch_dimensions
|
int64 repetido | lhs números de dimensão de lote |
rhs_batch_dimensions
|
int64 repetido | rhs números de dimensão de lote |
O DotGeneral executa a soma dos produtos acima das dimensões de contrato especificadas em
dimension_numbers
.
Os números de dimensão de contratação associados de lhs
e rhs
não precisam ser os mesmos, mas precisam ter os mesmos tamanhos de dimensão.
Exemplo com números de dimensão de contrato:
lhs = { {1.0, 2.0, 3.0},
{4.0, 5.0, 6.0} }
rhs = { {1.0, 1.0, 1.0},
{2.0, 2.0, 2.0} }
DotDimensionNumbers dnums;
dnums.add_lhs_contracting_dimensions(1);
dnums.add_rhs_contracting_dimensions(1);
DotGeneral(lhs, rhs, dnums) -> { {6.0, 12.0},
{15.0, 30.0} }
Os números de dimensão de lote associados de lhs
e rhs
precisam ter os mesmos
tamanhos de dimensão.
Exemplo com números de dimensão de lote (tamanho do lote 2, matrizes 2x2):
lhs = { { {1.0, 2.0},
{3.0, 4.0} },
{ {5.0, 6.0},
{7.0, 8.0} } }
rhs = { { {1.0, 0.0},
{0.0, 1.0} },
{ {1.0, 0.0},
{0.0, 1.0} } }
DotDimensionNumbers dnums;
dnums.add_lhs_contracting_dimensions(2);
dnums.add_rhs_contracting_dimensions(1);
dnums.add_lhs_batch_dimensions(0);
dnums.add_rhs_batch_dimensions(0);
DotGeneral(lhs, rhs, dnums) -> { { {1.0, 2.0},
{3.0, 4.0} },
{ {5.0, 6.0},
{7.0, 8.0} } }
Entrada | Saída | Semântica |
---|---|---|
[b0, m, k] dot [b0, k, n] |
[b0, m, n] | lote matmul |
[b0, b1, m, k] dot [b0, b1, k, n] |
[b0, b1, m, n] | lote matmul |
O número de dimensão resultante começa com a dimensão do lote, depois a dimensão lhs
sem contrato/sem lote e, por fim, a dimensão rhs
não contratante/sem lote.
DynamicSlice
Consulte também
XlaBuilder::DynamicSlice
.
DynamicSlice extrai uma submatriz da matriz de entrada na função 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ções em cada
dimensão: [início, início + tamanho). O formato de start_indices
precisa ser de classificação == 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 dimensional N do tipo T |
start_indices |
sequência de N XlaOp |
Lista de N números inteiros escalares contendo 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 start + size precisa ser menor ou igual ao tamanho da dimensão para evitar ajustar o tamanho da dimensão de módulo. |
Os índices de frações efetivas são calculados aplicando a seguinte
transformação para cada índice i
em [1, N)
antes de executar 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 sempre esteja dentro dos limites em relação à matriz de operandos. Se a fração estiver limitada antes da aplicação da transformação, a transformação 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
.
DynamicUpdateSlice gera um resultado que é o valor da matriz de entrada
operand
, com uma fração update
substituída em start_indices
.
O formato de update
determina o formato da submatriz do resultado que
é atualizado.
O formato de start_indices
precisa ser de 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 dimensional N do tipo T |
update |
XlaOp |
Matriz dimensional N do tipo T contendo a atualização da fração. Cada dimensão do formato de atualização precisa ser estritamente maior que zero, e start + update precisa ser menor ou igual ao tamanho do operando para 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 contendo os índices iniciais da fatia para cada dimensão. O valor precisa ser maior ou igual a zero. |
Os índices de frações efetivas são calculados aplicando a seguinte
transformação para cada índice i
em [1, N)
antes de executar 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 vinculada à matriz de operandos. Se a fração estiver limitada antes da aplicação da transformação, a transformação 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 com elementos
Consulte também
XlaBuilder::Add
.
Um conjunto de operações aritméticas binárias com elementos é permitido.
Op(lhs, rhs)
Em que Op
é Add
(adição), Sub
(subtração), Mul
(multiplicação), Div
(divisão), Rem
(restante), Max
(máximo), Min
(mínimo), LogicalAnd
(AND lógico) ou LogicalOr
(OR lógico).
Argumentos | Tipo | Semântica |
---|---|---|
lhs |
XlaOp |
operando do lado esquerdo: matriz do tipo T |
rhs |
XlaOp |
operando no lado direito: matriz do tipo T |
As formas dos argumentos precisam ser semelhantes ou compatíveis. Consulte a documentação sobre transmissão para entender o que significa que as formas são compatíveis. O resultado de uma operação tem uma forma que é resultado da transmissão das duas matrizes de entrada. Nessa variante, as operações entre matrizes de classificações diferentes não são compatíveis, a menos que um dos operandos seja um escalar.
Quando Op
é Rem
, o sinal do resultado é retirado do dividendo, e o
valor absoluto do resultado é sempre menor que o valor absoluto do divisor.
O estouro de divisão inteira (divisão/restante assinado/não assinado por zero ou divisão/restante
assinada de INT_SMIN
com -1
) produz um valor definido pela
implementação.
Há uma variante alternativa com suporte à transmissão de classificações diferentes para estas operações:
Op(lhs, rhs, broadcast_dimensions)
Em que Op
é igual ao exemplo acima. Essa variante da operação precisa ser usada
para operações aritméticas entre matrizes de classificações diferentes (como adicionar uma
matriz a um vetor).
O operando de classificação broadcast_dimensions
adicional é uma fração de números inteiros usado para
expandir a classificação do operando de classificação inferior até a classificação do operando de
classificação mais alta. broadcast_dimensions
mapeia as dimensões do formato de classificação inferior para as dimensões do formato de classificação superior. As dimensões não mapeadas da forma expandida são preenchidas com dimensões de tamanho um. A transmissão degeneração de dimensão
transmite as formas ao longo dessas dimensões degeneradas para equalizar as
formas de ambos os operandos. A semântica é descrita em detalhes na
página de transmissão.
Operações de comparação com elementos
Consulte também
XlaBuilder::Eq
.
Há suporte para um conjunto de operações de comparação binária de elemento padrão. Observe que a semântica padrão de comparação de ponto flutuante IEEE 754 se aplica 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), Lt
(menor que). Outro conjunto de operadores, EqTotalOrder, NeTotalOrder, GeTotalOrder,
GtTotalOrder, LeTotalOrder e LtTotalOrder, fornece as mesmas funcionalidades,
exceto que também dão suporte a um pedido total sobre os números de ponto
flutuante, aplicando -NaN < -Inf < -Finite < -0 < +0 < +Finite < +In.
Argumentos | Tipo | Semântica |
---|---|---|
lhs |
XlaOp |
operando do lado esquerdo: matriz do tipo T |
rhs |
XlaOp |
operando no lado direito: matriz do tipo T |
As formas dos argumentos precisam ser semelhantes ou compatíveis. Consulte a documentação sobre transmissão para entender o que significa que as formas são 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,
operações entre matrizes de classificações diferentes não são compatíveis, a menos que um
dos operandos seja um escalar.
Há uma variante alternativa com suporte à transmissão de classificações diferentes para estas operações:
Op(lhs, rhs, broadcast_dimensions)
Em que Op
é igual ao exemplo acima. Essa variante da operação precisa ser usada
para operações de comparação entre matrizes de classificações diferentes (como adicionar uma
matriz a um vetor).
O operando broadcast_dimensions
extra é uma fração 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 com elemento
O XlaBuilder é compatível com estas funções unárias com elementos:
Abs(operand)
absx -> |x|
por elemento.
Ceil(operand)
Ceil por elemento x -> ⌈x⌉
.
Cos(operand)
Cosseno por elemento x -> cos(x)
.
Exp(operand)
Exponencial natural por elemento x -> e^x
.
Floor(operand)
Valor mínimo do elemento x -> ⌊x⌋
.
Imag(operand)
Parte imaginária por elemento de uma forma complexa (ou real). x -> imag(x)
. Se o operando for um tipo de ponto flutuante, retorna 0.
IsFinite(operand)
Testa se cada elemento de operand
é finito,
ou seja, não é um 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
somente se o elemento de entrada correspondente for finito.
Log(operand)
Logaritmo natural por elemento x -> ln(x)
.
LogicalNot(operand)
A lógica por elemento não é x -> !(x)
.
Logistic(operand)
Computação da função logística por elemento x ->
logistic(x)
.
PopulationCount(operand)
Calcula o número de bits definidos em cada elemento de operand
.
Neg(operand)
Negação por elemento x -> -x
.
Real(operand)
Parte real com elementos de uma forma complexa (ou real).
x -> real(x)
. Se o operando for um tipo de ponto flutuante, retorna o mesmo valor.
Rsqrt(operand)
Recíproco elemento da operação de raiz quadrada
x -> 1.0 / sqrt(x)
.
Sign(operand)
Operação de sinal por elemento 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
.
Sqrt(operand)
Operação de raiz quadrada por elemento x -> sqrt(x)
.
Cbrt(operand)
Operação de raiz cúbica por elemento x -> cbrt(x)
.
Tanh(operand)
Tangente hiperbólica com elemento x -> tanh(x)
.
Round(operand)
Arredondamento por elemento, empata com zero.
RoundNearestEven(operand)
Arredondamento por elemento, vinculado ao par mais próximo.
Argumentos | Tipo | Semântica |
---|---|---|
operand |
XlaOp |
O operando para a função |
A função é aplicada a cada elemento na matriz operand
, resultando em uma
matriz com a mesma forma. operand
pode ser um escalar (classificação 0).
Fft
A operação FFT do XLA implementa as transformações de Fourier diretas e inversas para entradas/saídas reais e complexas. É possível usar FFTs multidimensionais em até três eixos.
Consulte também
XlaBuilder::Fft
.
Argumentos | Tipo | Semântica |
---|---|---|
operand |
XlaOp |
A matriz que estamos transformando Fourier. |
fft_type |
FftType |
Consulte a tabela abaixo. |
fft_length |
ArraySlice<int64> |
Os comprimentos de domínio de 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 |
Encaminhar FFT complexa para complexa. A forma não foi alterada. |
IFFT |
FFT inversa entre complexa. A forma não foi alterada. |
RFFT |
Encaminhar FFT real para complexa. 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 Nyquist. |
IRFFT |
FFT inversa entre real e complexa (ou seja, toma complexa, retorna real). A forma do eixo mais interno será 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 inverso das entradas 1 para fft_length[-1] // 2 + 1 . |
FFT multidimensional
Quando mais de um fft_length
é fornecido, isso equivale a aplicar uma
casada de operações FFT em cada um dos eixos mais internos. Observe que, para os casos reais > complexos e complexos > reais, a transformação mais interna do eixo é (efetivamente) executada primeiro (RFFT, por último para IRFFT), e é por isso que o eixo mais interno é aquele que muda de tamanho. Outras transformações de eixos serão
complexas> complexas.
Detalhes da implementação
A CPU FFT usa a TensorFFT do Eigen. A FFT da GPU usa cuFFT.
Gather
A operação de coleta do XLA une várias frações (cada fatia em um deslocamento do ambiente 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 informativa" abaixo.
gather(operand, start_indices, offset_dims, collapsed_slice_dims, slice_sizes, start_index_map)
Argumentos | Tipo | Semântica |
---|---|---|
operand |
XlaOp |
A matriz que estamos coletando. |
start_indices |
XlaOp |
Matriz contendo os índices iniciais das fatias reunidas. |
index_vector_dim |
int64 |
A dimensão em start_indices que "contém" os índices iniciais. Veja 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 a partir do operando. |
slice_sizes |
ArraySlice<int64> |
slice_sizes[i] são os limites para a fração na dimensão i . |
collapsed_slice_dims |
ArraySlice<int64> |
O conjunto de dimensões em cada fatia que foi recolhida. 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 |
Se os índices serão classificados pelo autor da chamada. |
Por conveniência, rotulamos as dimensões na matriz de saída que não estão em offset_dims
como batch_dims
.
O resultado é uma matriz da classificação 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
, vamos considerar implicitamente
start_indices
como tendo uma dimensão 1
à direita. Por exemplo, se start_indices
era da
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 junto à dimensão i
são calculados da seguinte maneira:
Se
i
estiver presente embatch_dims
(ou seja, for igual abatch_dims[k]
para algunsk
), vamos escolher os limites de dimensão correspondentes destart_indices.shape
, ignorandoindex_vector_dim
(ou seja, escolhastart_indices.shape.dims
[k
] sek
<index_vector_dim
estart_indices.shape.dims
[k
+1
] caso contrário).Se
i
estiver presente emoffset_dims
(ou seja, igual aoffset_dims
[k
] para algunsk
), vamos escolher o limite correspondente deslice_sizes
depois de contabilizarcollapsed_slice_dims
(ou seja, escolheradjusted_slice_sizes
[k
], em queadjusted_slice_sizes
éslice_sizes
, com os limites nos índicescollapsed_slice_dims
removidos).
Oficialmente, o índice de operando In
correspondente a um determinado índice de saída Out
é
calculado da seguinte maneira:
Permita que
G
= {Out
[k
] parak
embatch_dims
}. UseG
para separar um vetorS
de modo queS
[i
] =start_indices
[Combine(G
,i
)] em que Combine(A, b) insira b na posiçãoindex_vector_dim
em A. Observe que isso é bem definido mesmo queG
esteja vazio: seG
estiver vazio,S
=start_indices
.Crie um índice inicial,
S
in
, emoperand
usandoS
distribuindoS
comstart_index_map
. Mais precisamente:S
in
[start_index_map
[k
]] =S
[k
] sek
<start_index_map.size
.S
in
[_
] =0
caso contrário.
Crie um índice
O
in
emoperand
espalhando os índices nas dimensões de deslocamento emOut
de acordo com ocollapsed_slice_dims
definido. Mais precisamente:O
in
[remapped_offset_dims
(k
)] =Out
[offset_dims
[k
]] sek
<offset_dims.size
(remapped_offset_dims
é definido abaixo).O
in
[_
] =0
caso contrário.
In
éO
in
+S
in
, em que + é uma adição de 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, o XLA poderá presumir que start_indices
foram classificados (em ordem crescente de start_index_map
) pelo usuário. Se não forem, a semântica será definida pela implementação.
Descrição e exemplos informais
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 destart_indices
.Usamos
start_index_map
para mapear o índice inicial (cujo tamanho pode ser menor que opera.rank) para um índice inicial "completo" nooperand
.Dividamos dinamicamente uma fatia com o tamanho
slice_sizes
usando o índice inicial completo.Reunimos as dimensões
collapsed_slice_dims
para remodelar a fração. Como todas as dimensões das fatias recolhidas precisam ter um limite de 1, essa remodelação é sempre válida.Usamos as dimensões de deslocamento em
Out
para indexar essa fração e conseguir o elemento de entrada,E
, correspondente ao índice de saídaOut
.
index_vector_dim
está definido como start_indices.rank
– 1
em todos os exemplos
a seguir. Valores mais interessantes para index_vector_dim
não alteram a operação fundamentalmente, mas tornam a representação visual mais complicada.
Para ter uma intuição de como todos os itens acima se encaixam, vamos analisar um
exemplo que reúne cinco fatias da forma [8,6]
de uma matriz [16,11]
. A posição de uma fração na matriz [16,11]
pode ser representada como um vetor de índice da forma S64[2]
. Portanto, 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
,O
0
,O
1
], 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
,O
0
,O
1
] é o elemento na matriz de
entrada no índice [X
+O
0
,Y
+O
1
].
slice_sizes
é [8,6]
, que decide o intervalo de O0
e
O1
e, por sua vez, define os limites da fração.
Essa operação de coleta age como uma fração dinâmica de lote com G
como a dimensão do lote.
Os índices de coleta podem ser multidimensionais. Por exemplo, uma versão mais geral do exemplo acima usando uma matriz de "índices de coleta" da forma [4,5,2]
traduziria índices como este:
Novamente, isso age como uma fração dinâmica de lote G
0
e
G
1
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 das seguintes maneiras:
É possível configurar quais dimensões na forma de saída são as de deslocamento (no último exemplo, as que contêm
O
0
eO
1
). As dimensões do lote de saída (dimensões contendoG
0
,G
1
no último exemplo) são definidas como as dimensões de saída que não são deslocadas.O número de dimensões de deslocamento de saída explicitamente presentes no formato de saída pode ser menor 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 de1
. Como eles têm um tamanho de fatia de1
, o único índice válido para eles é0
e eliminá-los não introduz ambiguidade.A fração extraída da matriz "Gather Indices" ((
X
,Y
) no último exemplo) pode ter menos elementos do que a classificação da matriz de entrada. Um mapeamento explícito dita como o índice precisa ser expandido para ter a mesma classificação que a entrada.
No exemplo final, usamos (2) e (3) para implementar tf.gather_nd
:
G
0
e G
1
são usados para separar um índice inicial
da matriz de índices de coleta como de costume, mas o índice inicial tem apenas um
elemento, X
. Da mesma forma, há apenas um índice de deslocamento de saída com o valor
O
0
. No entanto, antes de serem usados como índices na matriz de entrada, eles são expandidos de acordo com "Reunir mapeamento de índice" (start_index_map
na descrição formal) e "Mapeamento de deslocamento" (remapped_offset_dims
na descrição formal) em [X
,0
] e [0
,O
0
] respectivamente, somando [X
,O
0
]. Em outras palavras, o índice de saída 0
, 0
, O
, o índice de saída 0
, 0
, 0
,G
O
G
G
G
1
1
GatherIndices
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 n 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 de dimensão n. |
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. |
Transmita o operando como resultado, com a dimensão dinâmica rastreada pelo compilador.
Os valores preenchidos 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
.
Índices em uma tupla com um valor constante de tempo de compilação.
O valor precisa ser uma constante de tempo de compilação para que a inferência de formas possa determinar o tipo do valor resultante.
Isso é análogo a std::get<int N>(t)
no 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
.
In-feed
Consulte também
XlaBuilder::Infeed
.
Infeed(shape)
Argumento | Tipo | Semântica |
---|---|---|
shape |
Shape |
Formato dos dados lidos da interface de 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 implícita de streaming de Infeed do dispositivo, interpretando os dados como a forma especificada e o layout, e retorna um XlaOp
dos dados. Várias operações de entrada são permitidas em um cálculo, mas é preciso haver um pedido total entre elas. Por
exemplo, dois feeds no código abaixo têm uma ordem total, porque há uma
dependência entre os loops "while".
result1 = while (condition, init = init_value) {
Infeed(shape)
}
result2 = while (condition, init = result1) {
Infeed(shape)
}
Formas de tupla aninhadas não são aceitas. Para uma forma de tupla vazia, a operação de Infeed é efetivamente um ambiente autônomo e prossegue sem ler nenhum dado da entrada 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 mantém valores a partir de zero e incrementando em um ao longo da dimensão especificada. Para tipos de
ponto flutuante, a matriz produzida é equivalente a ConvertElementType(Iota(...))
, em que
Iota
é do tipo integral e a conversão é para o tipo de ponto flutuante.
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 ]]
Mapeamento
Consulte também
XlaBuilder::Map
.
Map(operands..., computation)
Argumentos | Tipo | Semântica |
---|---|---|
operands |
sequência de N XlaOp s |
N matrizes dos tipos T0..T{N-1} |
computation |
XlaComputation |
computação do tipo T_0, T_1, .., T_{N + M -1} -> S com N parâmetros do tipo T e M do tipo arbitrário |
dimensions |
Matriz int64 |
matriz de dimensões de mapa |
Aplica uma função escalar nas matrizes operands
especificadas, produzindo uma matriz das mesmas dimensões, em que cada elemento é o resultado da função mapeada aplicada aos elementos correspondentes nas matrizes de entrada.
A função mapeada é um cálculo arbitrário com a restrição de que tem N entradas do tipo escalar T
e uma única saída com o tipo S
. A saída tem
as mesmas dimensões que os operandos, mas 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 de mover cálculos pela barreira.
Garante que todas as entradas sejam avaliadas antes de qualquer operador que dependa das saídas da barreira.
Acolchoamento
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 |
preenchimento nas duas bordas (baixa, alta) e entre os elementos de cada dimensão |
Expande a matriz de operand
especificada preenchendo a matriz e entre os elementos com o padding_value
especificado. padding_config
especifica a quantidade de padding da borda e do 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
na parte inferior (ao lado do índice 0) e na parte superior (próxima ao índice mais alto) de
cada dimensão, respectivamente. A quantidade de padding da borda pode ser negativa. O
valor absoluto do padding negativo indica o número de elementos que serão 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 preenchimento interno ocorre
logicamente antes do preenchimento da borda. Portanto, no caso do preenchimento da borda negativo, os elementos
são removidos do operando com preenchimento interno.
Essa operação será um ambiente autônomo se os pares de padding das bordas forem todos (0, 0) e os
valores do padding interno forem todos 0. A figura abaixo mostra exemplos de diferentes valores 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/recv |
Recebe dados da forma especificada de uma instrução Send
em outro cálculo que compartilha o mesmo identificador de canal. Retorna um
XlaOp para os dados recebidos.
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
a seguir para aguardar a conclusão da transferência
de dados. O contexto é uma tupla de {receive buffer (shape), request identificador
(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 conclusão da transferência de dados e retorna os dados recebidos.
Reduzir
Consulte também
XlaBuilder::Reduce
.
Aplica uma função de redução a uma ou mais matrizes em paralelo.
Reduce(operands..., init_values..., computation, dimensions)
Argumentos | Tipo | Semântica |
---|---|---|
operands |
Sequência de N XlaOp |
N matrizes dos tipos T_0, ..., T_{N-1} . |
init_values |
Sequência de N XlaOp |
N escalares dos tipos T_0, ..., T_{N-1} . |
computation |
XlaComputation |
de cálculo do tipo T_0, ..., T_{N-1}, T_0, ..., T_{N-1} -> Collate(T_0, ..., T_{N-1}) . |
dimensions |
Matriz int64 |
não ordenada de dimensões para reduzir. |
Em que:
- N precisa ser maior ou igual a 1.
- A computação precisa ser associativa "aproximadamente" (confira 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 elementosN
do tipoT
.
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
dimensões descritas abaixo.
Back-ends diferentes podem reassociar o cálculo de redução. Isso pode levar a diferenças numéricas, já que algumas funções de redução, como adição, não são associativas para flutuações. No entanto, se o intervalo dos dados for limitado, a adição de ponto flutuante será próxima o suficiente de ser associativa para os usos mais práticos.
Exemplos
Ao reduzir uma dimensão em uma única matriz 1D com os valores [10, 11,
12, 13]
, com a função de redução f
(computation
), isso seria calculado como:
f(10, f(11, f(12, f(init_value, 13)))
mas também há muitas outras possibilidades, como
f(init_value, f(f(10, f(init_value, 11)), f(f(init_value, 12), f(init_value, 13))))
Veja a seguir um exemplo aproximado de pseudocódigo de como a redução pode ser implementada, usando soma como o cálculo da 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]
Confira um exemplo de como reduzir uma matriz 2D (matriz). O formato tem classificação 2, dimensão 0 de tamanho 2 e dimensão 1 de tamanho 3:
Resultados da redução de dimensões 0 ou 1 com uma função "adicionar":
Observe que os dois resultados de redução são matrizes 1D. O diagrama mostra um como coluna e outro como linha apenas para conveniência visual.
Para um exemplo mais complexo, aqui está uma matriz 3D. A classificação dele é 3, a dimensão 0 de tamanho 4, a dimensão 1 de tamanho 2 e a dimensão 2 de tamanho 3. Para simplificar, os valores de 1 a 6 são replicados na dimensão 0.
Assim como no exemplo 2D, podemos reduzir apenas uma dimensão. Se reduzirmos a dimensão 0, por exemplo, obtemos uma matriz de classificação 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 obtemos uma matriz de classificação 2 em que todos os valores da dimensão 2 foram dobrados em um escalar:
| 6 15 |
| 6 15 |
| 6 15 |
| 6 15 |
A ordem relativa entre as dimensões restantes na entrada é preservada na saída, mas algumas dimensões podem receber novos números (já que a classificação muda).
Também podemos reduzir várias dimensões. A adição das dimensões 0 e 1 produz
a matriz 1D [20, 28, 36]
.
Reduzir a matriz 3D em todas as dimensões dela produz o 84
escalar.
Redução variadica
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
- ...
- Em execução no valor reduzido para o operando
- Valor de entrada para o primeiro operando
- ...
- Valor de entrada para o operador N'
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 1-D V = Float[N], K = Int[N]
e valores init
I_V = Float, I_K = Int
, o resultado f_(N-1)
da redução na única
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))
A aplicação dessa redução a uma matriz de valores e uma matriz de índices sequenciais (ou seja, iota), vai fazer uma coiteração nas matrizes e retornar uma tupla contendo 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 em um formato de menor precisão (como IEEE-FP16) e volta ao formato original. O número de bits de expoente e de mantissa no formato de menor precisão pode ser especificado arbitrariamente, embora nem todos os tamanhos de bits possam não ser compatíveis com 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 expoentes em formato de menor precisão |
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 representável com o número fornecido de bits de mantissa (usando a semântica
de "pares"). Todos os valores que excedem o intervalo especificado pelo número de
bits expoentes são fixados no infinito positivo ou negativo. Os valores NaN
são
mantidos, embora possam ser convertidos em valores NaN
canônicos.
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 mantissa zero) e
um número não negativo de bits de mantissa. O número de bits de expoente ou
mantissa pode exceder o valor correspondente para o tipo T
. A parte correspondente
da conversão é simplesmente um ambiente autônomo.
ReduceScatter
Consulte também
XlaBuilder::ReduceScatter
.
RedScatter é uma operação coletiva que faz um Allreduce e,
em seguida, dispersa o resultado dividindo-o em blocos shard_count
ao longo da
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 réplicas. |
computation |
XlaComputation |
Computação de 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 atual pode ser recuperado usandoReplicaId
). A ordem das réplicas em cada grupo determina a ordem em que o resultado de redução total será disperso.replica_groups
precisa estar vazio (nesse caso, todas as réplicas pertencem a um único grupo) ou conter o mesmo número de elementos que o número de réplicas. Quando há mais de um grupo de réplicas, todos eles precisam ser do mesmo tamanho. Por exemplo,replica_groups = {0, 2}, {1, 3}
executa a redução entre as réplicas0
e2
, e1
e3
, e depois espalha o resultado.shard_count
é o tamanho de cada grupo de réplicas. Precisamos disso nos casos em quereplica_groups
está vazio. Sereplica_groups
não estiver vazio,shard_count
precisará ser igual ao tamanho de cada grupo de réplicas.channel_id
é usado para comunicação entre módulos: apenas operaçõesreduce-scatter
com o mesmochannel_id
podem se comunicar entre si.
A forma de saída é a forma de entrada com scatter_dimension
tornado shard_count
vezes menor. Por exemplo, se houver duas réplicas e o
operando tiver os valores [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 ou uma tupla de N matrizes multidimensionais como saída. Cada matriz de saída tem o mesmo número de elementos que o
número de posições válidas da janela. Uma camada de pooling pode ser expressa como um
ReduceWindow
. Assim como no Reduce
, o computation
aplicado é
sempre transmitido ao 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 dos tipos T_0,..., T_{N-1} , cada uma representando a área de base em que a janela é colocada. |
init_values |
N XlaOps |
Os valores iniciais N 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 ser aplicada a elementos em cada janela de todos os operandos de entrada. |
window_dimensions |
ArraySlice<int64> |
matriz de números inteiros para valores de dimensão de janela |
window_strides |
ArraySlice<int64> |
matriz de números inteiros para valores de salto na janela |
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 de janela |
padding |
Padding |
tipo de padding para a janela (Padding::kSame, que preenche de modo a ter a mesma forma de saída da entrada se o stride for 1, ou Padding::kValid, que não usa padding e "para" a janela quando ela não cabe 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)
seráT
. - Se
N > 1
,Collate(T_0, ..., T_{N-1})
é uma tupla de elementosN
do tipo(T0,...T{N-1})
.
O código e a figura abaixo mostram um exemplo de como usar ReduceWindow
. A entrada é uma matriz de tamanho [4x6], e tanto window_dimensions quanto window_stride_dimensions são [2x3].
// Create a computation for the reduction (maximum).
XlaComputation max;
{
XlaBuilder builder(client_, "max");
auto y = builder.Parameter(0, ShapeUtil::MakeShape(F32, {}), "y");
auto x = builder.Parameter(1, ShapeUtil::MakeShape(F32, {}), "x");
builder.Max(y, x);
max = builder.Build().value();
}
// Create a ReduceWindow computation with the max reduction computation.
XlaBuilder builder(client_, "reduce_window_2x3");
auto shape = ShapeUtil::MakeShape(F32, {4, 6});
auto input = builder.Parameter(0, shape, "input");
builder.ReduceWindow(
input,
/*init_val=*/builder.ConstantLiteral(LiteralUtil::MinValue(F32)),
*max,
/*window_dimensions=*/{2, 3},
/*window_stride_dimensions=*/{2, 3},
Padding::kValid);
O salto igual a 1 em uma dimensão especifica que a posição de uma janela na dimensão está a um 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 stride diferentes. O padding é aplicado a cada dimensão da entrada, e os cálculos são os mesmos, como se a entrada tivesse vindo com as dimensões que ela tem após o preenchimento.
Para um exemplo de padding 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 salto 2
sobre a matriz de
entrada [10000, 1000, 100, 10, 1]
. O preenchimento kValid
calcula os mínimos em duas
janelas válidas: [10000, 1000, 100]
e [100, 10, 1]
, resultando na
saída [100, 1]
. O preenchimento de kSame
primeiro preenche a matriz para que a forma após
a janela de redução seja igual à entrada para o salto um, adicionando elementos
iniciais em ambos os lados, obtendo [MAX_VALUE, 10000, 1000, 100, 10, 1,
MAX_VALUE]
. A execução da janela de redução na 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 não ser determinista. Portanto, a função de redução não pode ser excessivamente
sensível à reassociaçã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
.
Reforma 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 de constante de tempo de compilação de índices de dimensão e um vetor de constante 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. Caso contrário, o padrão é {0, ..., rank - 1}
. A ordem das
dimensões em dimensions
vai da dimensão de variação mais lenta (mais importante) para
a dimensão de variação mais rápida (a mais secundária) no aninhamento de loop, que recolhe a
matriz de entrada em 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 das dimensões
do operando. Ao refinar a matriz recolhida para a matriz multidimensional definida por new_sizes
, as dimensões em new_sizes
são ordenadas da variação mais lenta (maior maior) até a variação mais rápida (menor).
Por exemplo, deixe v ser 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 remodelação pode transformar uma matriz de elemento único em uma 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 para reverter |
Inverte a ordem dos elementos na matriz operand
junto com o dimensions
especificado, gerando uma matriz de saída com a mesma forma. Cada elemento da
matriz de operandos em um índice multidimensional é armazenado na matriz de saída em um
índice transformado. O índice multidimensional é transformado pela inversão do índice em cada dimensão a ser invertida (por exemplo, se uma dimensão de tamanho N for uma das dimensões reversas, o índice i será transformado em N - 1 - i).
Um uso para a operação Rev
é reverter a matriz de peso de convolução ao longo das duas dimensões da janela durante o cálculo do gradiente em redes neurais.
RngNormal
Consulte também
XlaBuilder::RngNormal
.
Cria uma saída de uma determinada forma 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 do tipo |
shape |
Shape |
Forma de saída do tipo T |
RngUniform
Consulte também
XlaBuilder::RngUniform
.
Cria uma saída de uma determinada forma com números aleatórios gerados após a distribuição uniforme no intervalo \([a,b)\). Os parâmetros e o tipo de elemento de saída precisam ser booleanos, integrais ou 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 ser valores escalares. Se \(b <= a\) , o resultado for definido pela implementação.
RngUniform(a, b, shape)
Argumentos | Tipo | Semântica |
---|---|---|
a |
XlaOp |
Escalar do tipo T especificando o limite mínimo do intervalo |
b |
XlaOp |
Escalar do tipo T especificando o limite máximo 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 atual do número aleatório. O formato necessário e os valores válidos dependem do algoritmo usado.
A saída é garantida como uma função determinística do estado inicial, mas não há garantia de que ela seja determinista 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 os dados gerados. |
Valores disponíveis para algorithm
:
rng_default
: algoritmo específico de back-end com requisitos de forma específicos do back-end.rng_three_fry
: algoritmo PRNG baseado em contador da ThreeFry. A formainitial_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. O formatoinitial_state
éu64[3]
com valores arbitrários. Salmon et al. SC 2011. Números aleatórios paralelos: tão fáceis quanto 1, 2, 3.
Dispersão
A operação de dispersão de 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 dos tipos T_0, ..., T_N a serem espalhadas. |
scatter_indices |
XlaOp |
Matriz contendo os índices iniciais das frações que precisam ser distribuídas. |
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 espalhar operands[i] . |
update_computation |
XlaComputation |
Computação a ser usada para combinar os valores existentes na matriz de entrada e as atualizações durante o processo de dispersão. Esse cálculo 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 precisam ser inseridas no formato updates . |
scatter_dims_to_operand_dims |
ArraySlice<int64> |
Uma dimensão mapeia dos índices de dispersão para o espaço de índice de operandos. Essa matriz é interpretada como um mapeamento de i para scatter_dims_to_operand_dims[i] . Deve ser um para um e um total. |
indices_are_sorted |
bool |
Se os índices serão classificados pelo autor da chamada. |
unique_indices |
bool |
Se os índices têm a garantia de serem 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 elementosN
do tipoT
.
Se index_vector_dim
for igual a scatter_indices.rank
, vamos considerar implicitamente
scatter_indices
uma dimensão 1
à direita.
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
crescente.
Os argumentos de dispersão precisam seguir estas restrições:
Cada matriz
updates
precisa ter a classificaçãoupdate_window_dims.size + scatter_indices.rank - 1
.Os limites da dimensão
i
em cada matrizupdates
precisam estar em conformidade com o seguinte:- Se
i
estiver presente emupdate_window_dims
(ou seja, igual aupdate_window_dims
[k
] para algunsk
), o limite da dimensãoi
emupdates
não poderá exceder o limite correspondente deoperand
depois de contabilizarinserted_window_dims
(ou seja,adjusted_window_bounds
[k
], em queadjusted_window_bounds
contém os limites deoperand
com os limites nos índicesinserted_window_dims
removidos). - Se
i
estiver presente emupdate_scatter_dims
(ou seja, igual aupdate_scatter_dims
[k
] para algunsk
), o limite da dimensãoi
emupdates
precisará ser igual ao limite correspondente descatter_indices
, ignorandoindex_vector_dim
(ou seja,scatter_indices.shape.dims
[k
], sek
for menor queindex_vector_dim
escatter_indices.shape.dims
[k+1
] caso contrário).
- Se
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 deupdate_window_dims.size
einserted_window_dims.size
.scatter_dims_to_operand_dims.size
precisa ser igual ascatter_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:
- Permita que
G
= {U
[k
] parak
emupdate_scatter_dims
}. UseG
para procurar um vetor de índiceS
na matrizscatter_indices
de modo queS
[i
] =scatter_indices
[Combine(G
,i
)] em que Combine(A, b) insira b nas posiçõesindex_vector_dim
em A. - Crie um índice
S
in
emoperand
usandoS
espalhandoS
com o mapascatter_dims_to_operand_dims
. Mais formalmente:S
in
[scatter_dims_to_operand_dims
[k
]] =S
[k
] sek
<scatter_dims_to_operand_dims.size
.S
in
[_
] =0
caso contrário.
- Crie um índice
W
in
em cada matrizoperands
espalhando os índices emupdate_window_dims
emU
de acordo cominserted_window_dims
. Mais formalmente:W
in
[window_dims_to_operand_dims
(k
)] =U
[k
] sek
estiver emupdate_window_dims
, em quewindow_dims_to_operand_dims
é a função monotônica com domínio [0
,update_window_dims.size
) e intervalo [0
,operand.rank
) \inserted_window_dims
. Por exemplo, seupdate_window_dims.size
for4
,operand.rank
for6
einserted_window_dims
for {0
,2
}, entãowindow_dims_to_operand_dims
será {0
→1
,1
→3
,2
→4
,3
→5
}.W
in
[_
] =0
caso contrário.
I
éW
in
+S
in
, em que + é uma adição de elemento.
Em resumo, a operação de dispersão pode ser definida da seguinte maneira:
- Inicialize
output
comoperands
, ou seja, para todos os índicesJ
, para todos os índicesO
na matrizoperands
[J
]:
output
[J
][O
] =operands
[J
][O
] - Para cada índice
U
na matrizupdates
[J
] e o índice correspondenteO
na matrizoperand
[J
], seO
for um índice válido paraoutput
:
(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 é determinista. Portanto, quando vários
índices em updates
se referirem ao mesmo índice em operands
, o valor
correspondente em output
será não determinístico.
O primeiro parâmetro transmitido para update_computation
será sempre o valor atual da matriz output
e o segundo parâmetro sempre será o valor da matriz updates
. Isso é importante
especificamente nos casos em que update_computation
não é comutativo.
Se indices_are_sorted
for definido como verdadeiro, o XLA poderá presumir que start_indices
foram classificados (em ordem crescente de start_index_map
) pelo usuário. Se não forem, a semântica será definida pela implementação.
Se unique_indices
for definido como verdadeiro, o XLA poderá presumir que todos os elementos
dispersos são exclusivos. Portanto, o XLA pode usar operações não atômicas. Se
unique_indices
for definido como verdadeiro e os índices que estão sendo espalhados não forem
exclusivos, a semântica será implementada.
informalmente, a op de dispersão pode ser considerada um inverso da operação de coleta, ou seja, a op de dispersão atualiza os elementos na entrada que são extraídos pela operação de coleta correspondente.
Para uma descrição informal detalhada e exemplos, consulte a seção "Descrição informativa" em Gather
.
Selecione
Consulte também
XlaBuilder::Select
.
Constrói uma matriz de saída com elementos de duas matrizes de entrada, com base nos valores de uma matriz de predicado.
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 de
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 será
tirado de on_true
se o valor de P
for true
e de on_false
se o
valor de P
for false
. Como uma forma restrita de transmissão,
pred
pode ser um escalar do tipo PRED
. Nesse caso, a matriz de saída é totalmente extraída 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};
Há suporte para seleções entre tuplas. As tuplas são consideradas tipos escalares
para essa finalidade. Se on_true
e on_false
forem tuplas (que precisam ter a mesma forma), pred
terá que ser um escalar do tipo PRED
.
SelectAndScatter
Consulte também
XlaBuilder::SelectAndScatter
.
Essa operação pode ser considerada uma operação composta que primeiro calcula
ReduceWindow
na matriz operand
para selecionar um elemento de cada janela e,
em seguida, dispersa a matriz source
para os índices dos elementos selecionados para
criar uma matriz de saída com a mesma forma da matriz de operandos. A função binária select
é usada para selecionar um elemento de cada janela aplicando-a 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
retornará true
se o primeiro parâmetro for
selecionado e false
se o segundo parâmetro for selecionado. A
função precisa conter a transitividade (ou seja, se select(a, b)
e select(b, c)
forem
true
, então select(a, c)
também será true
) para que o elemento selecionado
não dependa da ordem dos elementos percorridos em uma determinada janela.
A função scatter
é aplicada a cada índice selecionado na matriz de saída. Ele usa dois parâmetros escalares:
- Valor atual no índice selecionado na matriz de saída
- O valor de dispersão de
source
que se aplica ao índice selecionado
Ele combina os dois parâmetros e retorna um valor escalar que é usado para atualizar o valor no índice selecionado na matriz de saída. Inicialmente, todos os índices da
matriz de saída são definidos como init_value
.
A matriz de saída tem a mesma forma que a matriz operand
, e a matriz source
precisa ter a mesma forma do resultado da aplicação de uma operação ReduceWindow
na matriz operand
. SelectAndScatter
pode ser usado para fazer retropropagação dos valores do gradiente para uma camada de pooling 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 |
computação binária do tipo T, T -> PRED , para aplicar a todos os elementos em cada janela; retorna true se o primeiro parâmetro for selecionado e retorna false se o segundo parâmetro for selecionado |
window_dimensions |
ArraySlice<int64> |
matriz de números inteiros para valores de dimensão de janela |
window_strides |
ArraySlice<int64> |
matriz de números inteiros para valores de salto na janela |
padding |
Padding |
tipo de padding da janela (Padding::kSame ou Padding::kValid) |
source |
XlaOp |
matriz do tipo T com os valores a serem distribuídos |
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 seu elemento de destino |
A figura abaixo mostra exemplos de uso de SelectAndScatter
, com a função select
calculando o valor máximo entre os parâmetros. Quando as
janelas se sobrepõem, como na figura (2) abaixo, um índice da matriz operand
pode
ser selecionado várias vezes em diferentes janelas. Na figura, o elemento de
valor 9 é selecionado pelas duas janelas superiores (azul e vermelho), 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 não ser determinista. Portanto, a função scatter
não pode ser excessivamente
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/recv |
Envia os dados de operando fornecidos para uma instrução Recv
em outro cálculo
que compartilha o mesmo identificador de canal. Não retorna dados.
Semelhante à operação Recv
, a API do 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 do 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 {operand (shape), request identificador
(U32)} e só pode ser usado por uma instrução SendDone
.
SendDone(HloInstruction context)
Em 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 deSend
Send
acontece antes deRecvDone
Recv
acontece antes deRecvDone
Send
acontece antes deSendDone
Quando os compiladores de back-end geram uma programação linear para cada computação que se comunica por instruções de canal, não pode haver ciclos entre os 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 fatia.
Slice(operand, start_indices, limit_indices, strides)
Argumentos | Tipo | Semântica |
---|---|---|
operand |
XlaOp |
Matriz dimensional N do tipo T |
start_indices |
ArraySlice<int64> |
Lista de N números inteiros contendo os índices iniciais da fração de cada dimensão. Os valores devem ser maiores ou iguais a zero. |
limit_indices |
ArraySlice<int64> |
Lista de N números inteiros contendo os índices finais (exclusivos) da fatia de cada dimensão. Cada valor precisa ser maior ou igual ao respectivo valor de start_indices para a dimensão e menor ou igual ao tamanho da dimensão. |
strides |
ArraySlice<int64> |
Lista de N números inteiros que decide o salto de entrada da fração. A fração escolhe cada elemento 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} }
Ordenar
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 |
O cálculo do comparador a ser usado. |
dimension |
int64 |
A dimensão com a qual classificar. |
is_stable |
bool |
Define 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á realizar uma comparação "menor que". Oficialmente, depois que a matriz é classificada, ela é mantida para todas as posições de índice
i, j
comi < j
que sãocomparator(value[i], value[j]) = comparator(value[j], value[i]) = false
oucomparator(value[i], value[j]) = true
.Se o operando tiver uma classificação mais alta, ele será classificado junto com a dimensão fornecida. Por exemplo, para um tensor de classificação 2 (uma matriz), um valor de dimensão de
0
classificará de maneira independente cada coluna, e um valor de dimensão de1
classificará de maneira independente cada linha. Se nenhum número de dimensão for informado, a última dimensão será escolhida por padrão. Para a dimensão classificada, a mesma ordem de classificação se aplica 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
ej
precisam ser trocados, o comparador é chamado com parâmetros escalares2 * n
, em que o parâmetro2 * k
corresponde ao valor na posiçãoi
do operandok-th
, e o parâmetro2 * k + 1
corresponde ao valor na posiçãoj
do operandok-th
. Normalmente, o comparador compara os parâmetros2 * k
e2 * k + 1
entre si e possivelmente usa outros pares de parâmetros como desempadores.O resultado é uma tupla que consiste nos operandos em ordem de classificação (ao longo da dimensão fornecida, conforme mostrado acima). O operando
i-th
da tupla corresponde ao operandoi-th
de Sort.
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 classificação será estável, ou seja, se
houver elementos considerados iguais pelo comparador, a
ordem relativa dos valores iguais será preservada. Dois elementos, e1
e e2
, serã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 fornecida, portanto,
∀ i . 0 ≤ i < rank ⇒ input_dimensions[permutation[i]] = output_dimensions[i]
.
É o mesmo que Reshape(operando, permutação, Permute(permute, operand.shape.dimensions)).
TriangularSolve
Consulte também
XlaBuilder::TriangularSolve
.
Resolve sistemas de equações lineares com matrizes de coeficientes triangulares superiores ou inferiores por substituição para frente ou para trás. Transmitindo ao longo de 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
, ou 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 classificação > 2 do mesmo tipo com a forma [..., M, K] se left_side for verdadeiro. Caso contrário, [..., K, M] . |
left_side |
bool |
indica se deve resolver um sistema no formato op(a) * x = b (true ) ou x * op(a) = b (false ). |
lower |
bool |
se o triângulo superior ou inferior de a será usado. |
unit_diagonal |
bool |
Se true , os elementos diagonais de a serão considerados 1 e não acessados. |
transpose_a |
Transpose |
usar a como está, transpor a função ou receber a transposição conjugada dela. |
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, e 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, eles serão tratados como lotes de matrizes, em que todas, exceto as duas dimensões secundárias, são dimensões em lote. a
e b
precisam ter dimensões de lote iguais.
Tupla
Consulte também
XlaBuilder::Tuple
.
Uma tupla contendo um número variável de identificadores de dados, cada um com o próprio formato.
Isso é análogo a std::tuple
no 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 desconstruídas (acessadas) com a 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 encerramento 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 a body
sequencialmente até que o condition
falhe. Isso é semelhante a
uma repetição "while" comum em muitas outras linguagens, exceto pelas diferenças e
restrições listadas abaixo.
- Um nó
While
retorna um valor do tipoT
, que é o resultado da última execução dobody
. - A forma do tipo
T
é determinada estaticamente 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 são atualizados automaticamente para o novo resultado de body
em cada iteração subsequente.
Um caso de uso principal 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 o cálculo. 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. Por 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};
}